Repository: kubernetes-sigs/kueue

Issue:
Kubeflow Trainer v2 integration follow ups
<!-- Please only use this template for submitting clean up requests -->

**What would you like to be cleaned**:

This is to address remaining comments to https://github.com/kubernetes-sigs/kueue/pull/6597

1. Don't lookup the JobSet object in the GenericJob implementation
2. Don't use the Background context, instead use the "request" context from PodSets() etc.


**Why is this needed**:
1 because:
a) JobSet should be an implementation detail of the Trainer which might be replaced in the future
b) it creates tricky race conditions to reason about, because the JobSet may not be immediately created
c) it couples the JobSet API version used for Trainer with the version we integrate with

Instead, helpers exposed by the Trainer should be used, but they require releasing 2.0.1.

2. becuase:
1. background context cannot be cancelled, so the goroutine leaks / remains running even if the main Kueue goroutine is terminated
2. the "request" context may have more data to log, like the workload name etc.

Discussion:
@mimowo: /assign @kaisoz 
Who I think is already looking into the follow ups.
cc @tenzen-y @andreyvelich @astefanutti @kannon92: Also can we add a page in our site for trainer?

cc @andreyvelich @kaisoz@mimowo: For sure we need docs, but we can add docs post code release, potentially. 

So, I would prefer to focus on addressing the code follow ups. @izturn: Just tried the current code and it doesn’t seem to work in **TAS** mode. Looks like the Pods don’t get the proper podset/workload labels and annotations when they start up(at `RunWithPodSetsInfo`). Because of that, the ungater can’t find the matching Pods(at `podsForPodSet`), so the scheduling gate never gets removed and the pods just stays stuck in **Pending**. @kaisoz @mimowo: Yes, this is the current "known limitation" due to the fact that Trainer v2 does not allow to update the podSetOverrides with labels & annotations, see open issue: https://github.com/kubernetes-sigs/kueue/issues/6865

It would be great to work on resoling this limitation.

cc @mwysokin @mimowo: Actually, there might be a dirty workaround inside Kueue. We could access the JobSet object and inject the labels & annotations before unsuspending in the Resume phase. 

If we have a PR, I could support it as a "temporary fix" while waiting for Trainer v2 fixes.@mimowo: cc @kaisoz @andreyvelich @tenzen-y @kaisoz: > Actually, there might be a dirty workaround inside Kueue. We could access the JobSet object and inject the labels & annotations before unsuspending in the Resume phase. 
> 
> If we have a PR, I could support it as a "temporary fix" while waiting for Trainer v2 fixes.

Unfortunately that won't work @mimowo because the trainjob controller reconciles the jobset and removes the labels and annotations (I've tried this already)@mimowo: Got it, what about mutating the Trainer's Runtime object before resuming@andreyvelich: @mimowo @kaisoz Can we configure those labels/annotations in the Runtime before TrainJob is created?
Or labels must be set by the Kueue controller ?@izturn: the `podsForPodSet`'s filter is the podset & workload's name, so i guess they can't be configured?

> [@mimowo](https://github.com/mimowo) [@kaisoz](https://github.com/kaisoz) Can we configure those labels/annotations in the Runtime before TrainJob is created? Or labels must be set by the Kueue controller ?

@mimowo: > Can we configure those labels/annotations in the Runtime before TrainJob is created?

before created I don't think, because Kueue may admit / evict multiple times during lifetime.

We could however "modify" the TainerRuntime object, unless it is explicitly prohibited by a webhook?@andreyvelich: >before created I don't think, because Kueue may admit / evict multiple times during lifetime.

Does Kueue modify the TAS labels during admission or eviction ? @mimowo: > Does Kueue modify the TAS labels during admission or eviction ?

Both@mimowo: @izturn one more thought - can you check using "main" Kueue. I suspect that maybe there is a difference after we stopped injecting the "tas" label. Still, we will not support the "implicit" mode, but I suspect that "required" and "preferred" modes have a chance to work.@tenzen-y: Additionally, I would like to narrow down the below annotation usage by renaming to `kueue.x-k8s.io/trainjob-override-idx` so that we can avoid accidentally introduce the annotation to different places.

https://github.com/kubernetes-sigs/kueue/blob/4c66b01c67e972c7250eeeb8ffeb87ccf961054a/pkg/controller/jobs/kubeflow/trainjob/trainjob_controller.go#L53
@mimowo: Ultimately I think the podSetOverrides should have the 'manager' field, then the annotation would not be needed. I discussed this with  Tomas and we came up with the annotation idea as a workaround, so that Kueue can know which overrides are done by Kueue and which by user.@kaisoz: > Got it, what about mutating the Trainer's Runtime object before resuming

I think that should work! Although it'd make both the `RunWithPodsetsInfo` and `RestorePodSetsInfo` a bit more complicated. Do you want me to try that out or better wait for Trainer PR to sink in?@kaisoz: > Additionally, I would like to narrow down the below annotation usage by renaming to `kueue.x-k8s.io/trainjob-override-idx` so that we can avoid accidentally introduce the annotation to different places.
> 
> [kueue/pkg/controller/jobs/kubeflow/trainjob/trainjob_controller.go](https://github.com/kubernetes-sigs/kueue/blob/4c66b01c67e972c7250eeeb8ffeb87ccf961054a/pkg/controller/jobs/kubeflow/trainjob/trainjob_controller.go#L53)
> 
> Line 53 in [4c66b01](/kubernetes-sigs/kueue/commit/4c66b01c67e972c7250eeeb8ffeb87ccf961054a)
> 
>  FirstOverrideIdx = "kueue.x-k8s.io/override-idx"

As @mimowo said, this is a workaround,but anyway I can rename the annotation 👍🏻 @izturn: > [@izturn](https://github.com/izturn) one more thought - can you check using "main" Kueue. I suspect that maybe there is a difference after we stopped injecting the "tas" label. Still, we will not support the "implicit" mode, but I suspect that "required" and "preferred" modes have a chance to work.

Sure, I can try again. But tbh, as a beginner, I don’t really understand what you mean by the “main” kueue things. I’m not sure how to do that. Could you walk me through it with a bit more detail? Thanks!@tenzen-y: > Ultimately I think the podSetOverrides should have the 'manager' field, then the annotation would not be needed. I discussed this with Tomas and we came up with the annotation idea as a workaround, so that Kueue can know which overrides are done by Kueue and which by user.

We probably need to generize the case for manager, then propose it to KF Trainer and it will not contain in the next Kueue release (v0.14).
Meanwhile, I would like to use `kueue.x-k8s.io/trainjob-override-idx` or `alpha.kueue.x-k8s.io/trainjob-override-idx` annotation.
I'm leaning towards annotation with "alpha" prefix since it should be a temporary annotation.
@mimowo: Regarding the alpha prefix I think there used to be this guidance in the past, but the API reviewers moved away from putting maturity level into the label or annotation name. I added the backOffLinitPerIndex feature which introduces annotation and I wasn't asked to add the prefix. Also I see other Alpha annotations in 1.34 not adding this https://kubernetes.io/docs/reference/labels-annotations-taints/?utm_source=chatgpt.com#applyset-kubernetes-io-additional-namespaces
so the maturity level is emphasized in the docs, but not in the annotation name.

I think if we just add to comment: This is alpha level annotation it should be enough. We could also call it deprecated already which maybe is better maturity level. I will open an issue in trainer.@mimowo: As part of this cleanup task we can unexport the const field @tenzen-y: > Regarding the alpha prefix I think there used to be this guidance in the past, but the API reviewers moved away from putting maturity level into the label or annotation name. I added the backOffLinitPerIndex feature which introduces annotation and I wasn't asked to add the prefix. Also I see other Alpha annotations in 1.34 not adding this https://kubernetes.io/docs/reference/labels-annotations-taints/?utm_source=chatgpt.com#applyset-kubernetes-io-additional-namespaces so the maturity level is emphasized in the docs, but not in the annotation name.
> 
> I think if we just add to comment: This is alpha level annotation it should be enough. We could also call it deprecated already which maybe is better maturity level. I will open an issue in trainer.

That makes sense. In that case, let's use `kueue.x-k8s.io/trainjob-override-idx`.@mimowo: > > [@izturn](https://github.com/izturn) one more thought - can you check using "main" Kueue. I suspect that maybe there is a difference after we stopped injecting the "tas" label. Still, we will not support the "implicit" mode, but I suspect that "required" and "preferred" modes have a chance to work.
> 
> Sure, I can try again. But tbh, as a beginner, I don’t really understand what you mean by the “main” kueue things. I’m not sure how to do that. Could you walk me through it with a bit more detail? Thanks!


Sure, by "main" I meant the latest build of Kueue from the main branch. You can install for example following https://kueue.sigs.k8s.io/docs/installation/#install-the-latest-development-version


However, I think it will not work anyway on the second thought, because we would not be able to inject the "workload" annotation :/ @izturn: > > > [@izturn](https://github.com/izturn) one more thought - can you check using "main" Kueue. I suspect that maybe there is a difference after we stopped injecting the "tas" label. Still, we will not support the "implicit" mode, but I suspect that "required" and "preferred" modes have a chance to work.
> > 
> > 
> > Sure, I can try again. But tbh, as a beginner, I don’t really understand what you mean by the “main” kueue things. I’m not sure how to do that. Could you walk me through it with a bit more detail? Thanks!
> 
> Sure, by "main" I meant the latest build of Kueue from the main branch. You can install for example following https://kueue.sigs.k8s.io/docs/installation/#install-the-latest-development-version
> 
> However, I think it will not work anyway on the second thought, because we would not be able to inject the "workload" annotation :/

Oh, got it. My reply above was actually based on testing with the main branch.@kaisoz: /reopen 

closed by mistake, `2` should still be addressed@k8s-ci-robot: @kaisoz: Reopened this issue.

<details>

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

>/reopen 
>
>closed by mistake, `2` should still be addressed


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>@tenzen-y: @mimowo @kaisoz Additionally, I would request 2 things as I mentioned in https://github.com/kubernetes-sigs/kueue/pull/7030#pullrequestreview-3283582303.
@kaisoz: > [@mimowo](https://github.com/mimowo) [@kaisoz](https://github.com/kaisoz) Additionally, I would request 2 things as I mentioned in [#7030 (review)](https://github.com/kubernetes-sigs/kueue/pull/7030#pullrequestreview-3283582303).

Hi @tenzen-y , just saw this, I'll take care of it 😊 
@tenzen-y: > > [@mimowo](https://github.com/mimowo) [@kaisoz](https://github.com/kaisoz) Additionally, I would request 2 things as I mentioned in [#7030 (review)](https://github.com/kubernetes-sigs/kueue/pull/7030#pullrequestreview-3283582303).
> 
> Hi [@tenzen-y](https://github.com/tenzen-y) , just saw this, I'll take care of it 😊

Thank you!@tenzen-y: x-ref: https://github.com/kubernetes-sigs/kueue/pull/7081#discussion_r2391731893@kaisoz: > x-ref: [#7081 (comment)](https://github.com/kubernetes-sigs/kueue/pull/7081#discussion_r2391731893)

PR ready @tenzen-y  https://github.com/kubernetes-sigs/kueue/pull/7110
Thanks!

Solution patch:
diff --git a/go.mod b/go.mod
index ea4f2ec3a..b19fed1eb 100644
--- a/go.mod
+++ b/go.mod
@@ -9,7 +9,7 @@ require (
 	github.com/google/go-cmp v0.7.0
 	github.com/json-iterator/go v1.1.12
 	github.com/kubeflow/mpi-operator v0.6.0
-	github.com/kubeflow/trainer/v2 v2.0.1-0.20250924122344-56be1f8163fb
+	github.com/kubeflow/trainer/v2 v2.0.1
 	github.com/kubeflow/training-operator v1.9.3
 	github.com/onsi/ginkgo/v2 v2.25.3
 	github.com/onsi/gomega v1.38.2
diff --git a/go.sum b/go.sum
index 262517328..c19dbce17 100644
--- a/go.sum
+++ b/go.sum
@@ -124,8 +124,8 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
 github.com/kubeflow/mpi-operator v0.6.0 h1:RzzG03oyQIiHe/SH2k19N+Hp4QkxEsG2l9G1T7US+ys=
 github.com/kubeflow/mpi-operator v0.6.0/go.mod h1:lu3cj3ESq3SdS29nbHC26TuXiTYYQYTHJ6fcy5Xv20c=
-github.com/kubeflow/trainer/v2 v2.0.1-0.20250924122344-56be1f8163fb h1:EN3ut+K1x6ZevkNT56r6zgbKFYZp9aNkXGiEZjcUR94=
-github.com/kubeflow/trainer/v2 v2.0.1-0.20250924122344-56be1f8163fb/go.mod h1:AH97YWUun+pLQlQHyRIlYcXw+YUuuQhnVWVYTbzGvgo=
+github.com/kubeflow/trainer/v2 v2.0.1 h1:4bMZdFIpQBHjcUene6P+MqbY5v7ccyLUIKkefOhUFG8=
+github.com/kubeflow/trainer/v2 v2.0.1/go.mod h1:AH97YWUun+pLQlQHyRIlYcXw+YUuuQhnVWVYTbzGvgo=
 github.com/kubeflow/training-operator v1.9.3 h1:aaSHqskOtCY8Dn8sVd+l9p5XAczW6O0nYMPjVLAw/lc=
 github.com/kubeflow/training-operator v1.9.3/go.mod h1:6zI0hgeCOheiW5Z12IcVkKBuSWm714fz8mUvYu4Fid4=
 github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc=
diff --git a/pkg/controller/jobframework/tas_validation.go b/pkg/controller/jobframework/tas_validation.go
index 6cd0f2176..733598976 100644
--- a/pkg/controller/jobframework/tas_validation.go
+++ b/pkg/controller/jobframework/tas_validation.go
@@ -34,6 +34,7 @@ func ValidateTASPodSetRequest(replicaPath *field.Path, replicaMetadata *metav1.O
 	_, unconstrainedFound := replicaMetadata.Annotations[kueuebeta.PodSetUnconstrainedTopologyAnnotation]
 	sliceRequiredValue, sliceRequiredFound := replicaMetadata.Annotations[kueuebeta.PodSetSliceRequiredTopologyAnnotation]
 	_, sliceSizeFound := replicaMetadata.Annotations[kueuebeta.PodSetSliceSizeAnnotation]
+	podSetGroupNameValue, podSetGroupNameFound := replicaMetadata.Annotations[kueuebeta.PodSetGroupName]
 
 	// validate no more than 1 annotation
 	asInt := func(b bool) int {
@@ -65,7 +66,6 @@ func ValidateTASPodSetRequest(replicaPath *field.Path, replicaMetadata *metav1.O
 	}
 
 	// validate PodSetGroupName annotation
-	podSetGroupNameValue, podSetGroupNameFound := replicaMetadata.Annotations[kueuebeta.PodSetGroupName]
 	if podSetGroupNameFound {
 		allErrs = append(allErrs, validatePodSetGroupNameAnnotation(podSetGroupNameValue, annotationsPath.Key(kueuebeta.PodSetGroupName))...)
 	}
@@ -82,10 +82,18 @@ func ValidateTASPodSetRequest(replicaPath *field.Path, replicaMetadata *metav1.O
 			allErrs = append(allErrs, field.Required(annotationsPath.Key(kueuebeta.PodSetSliceSizeAnnotation), "slice size is required if slice topology is requested"))
 		}
 	}
-
 	if !sliceRequiredFound && sliceSizeFound {
 		allErrs = append(allErrs, field.Forbidden(annotationsPath.Key(kueuebeta.PodSetSliceSizeAnnotation), fmt.Sprintf("cannot be set when '%s' is not present", kueuebeta.PodSetSliceRequiredTopologyAnnotation)))
 	}
+	if podSetGroupNameFound {
+		if sliceSizeFound {
+			allErrs = append(allErrs, field.Forbidden(annotationsPath.Key(kueuebeta.PodSetSliceSizeAnnotation), fmt.Sprintf("cannot be set when '%s' is present", kueuebeta.PodSetGroupName)))
+		}
+
+		if sliceRequiredFound {
+			allErrs = append(allErrs, field.Forbidden(annotationsPath.Key(kueuebeta.PodSetSliceRequiredTopologyAnnotation), fmt.Sprintf("cannot be set when '%s' is present", kueuebeta.PodSetGroupName)))
+		}
+	}
 
 	return allErrs
 }
diff --git a/vendor/modules.txt b/vendor/modules.txt
index 9d22e8255..fa0583ec5 100644
--- a/vendor/modules.txt
+++ b/vendor/modules.txt
@@ -190,7 +190,7 @@ github.com/json-iterator/go
 # github.com/kubeflow/mpi-operator v0.6.0
 ## explicit; go 1.23
 github.com/kubeflow/mpi-operator/pkg/apis/kubeflow/v2beta1
-# github.com/kubeflow/trainer/v2 v2.0.1-0.20250924122344-56be1f8163fb
+# github.com/kubeflow/trainer/v2 v2.0.1
 ## explicit; go 1.24.0
 github.com/kubeflow/trainer/v2/pkg/apis/trainer/v1alpha1
 github.com/kubeflow/trainer/v2/pkg/apply



Test patch:
diff --git a/pkg/controller/jobs/leaderworkerset/leaderworkerset_webhook_test.go b/pkg/controller/jobs/leaderworkerset/leaderworkerset_webhook_test.go
index 56065e563..d10a3a911 100644
--- a/pkg/controller/jobs/leaderworkerset/leaderworkerset_webhook_test.go
+++ b/pkg/controller/jobs/leaderworkerset/leaderworkerset_webhook_test.go
@@ -259,6 +259,41 @@ func TestValidateCreate(t *testing.T) {
 			}.ToAggregate(),
 			topologyAwareScheduling: true,
 		},
+		"invalid slice topology request - slicing requested together with grouping": {
+			lws: testingleaderworkerset.MakeLeaderWorkerSet("test-lws", "").
+				Queue("test-queue").
+				LeaderTemplate(corev1.PodTemplateSpec{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							kueuealpha.PodSetGroupName:                       "leadername1",
+							kueuealpha.PodSetSliceRequiredTopologyAnnotation: "cloud.com/block",
+							kueuealpha.PodSetSliceSizeAnnotation:             "1",
+						},
+					},
+				}).
+				Size(4).
+				WorkerTemplate(corev1.PodTemplateSpec{
+					ObjectMeta: metav1.ObjectMeta{
+						Annotations: map[string]string{
+							kueuealpha.PodSetGroupName:                       "workername1",
+							kueuealpha.PodSetSliceRequiredTopologyAnnotation: "cloud.com/block",
+							kueuealpha.PodSetSliceSizeAnnotation:             "1",
+						},
+					},
+				}).
+				Obj(),
+			wantErr: field.ErrorList{
+				field.Forbidden(field.NewPath("spec.leaderWorkerTemplate.leaderTemplate.metadata.annotations").
+					Key("kueue.x-k8s.io/podset-slice-size"), "cannot be set when 'kueue.x-k8s.io/podset-group-name' is present"),
+				field.Forbidden(field.NewPath("spec.leaderWorkerTemplate.leaderTemplate.metadata.annotations").
+					Key("kueue.x-k8s.io/podset-slice-required-topology"), "cannot be set when 'kueue.x-k8s.io/podset-group-name' is present"),
+				field.Forbidden(field.NewPath("spec.leaderWorkerTemplate.workerTemplate.metadata.annotations").
+					Key("kueue.x-k8s.io/podset-slice-size"), "cannot be set when 'kueue.x-k8s.io/podset-group-name' is present"),
+				field.Forbidden(field.NewPath("spec.leaderWorkerTemplate.workerTemplate.metadata.annotations").
+					Key("kueue.x-k8s.io/podset-slice-required-topology"), "cannot be set when 'kueue.x-k8s.io/podset-group-name' is present"),
+			}.ToAggregate(),
+			topologyAwareScheduling: true,
+		},
 		"valid PodSet group name request": {
 			lws: testingleaderworkerset.MakeLeaderWorkerSet("test-lws", "").
 				Queue("test-queue").
diff --git a/pkg/util/testingjobs/trainjob/wrappers.go b/pkg/util/testingjobs/trainjob/wrappers.go
index e761c7ffd..87484567e 100644
--- a/pkg/util/testingjobs/trainjob/wrappers.go
+++ b/pkg/util/testingjobs/trainjob/wrappers.go
@@ -157,3 +157,18 @@ func MakeClusterTrainingRuntime(name string, jobsetSpec jobsetapi.JobSetSpec) *k
 		},
 	}
 }
+
+// MakeTrainingRuntime creates a TrainingRuntime with the jobsetSpec provided
+func MakeTrainingRuntime(name, ns string, jobsetSpec jobsetapi.JobSetSpec) *kftrainerapi.TrainingRuntime {
+	return &kftrainerapi.TrainingRuntime{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      name,
+			Namespace: ns,
+		},
+		Spec: kftrainerapi.TrainingRuntimeSpec{
+			Template: kftrainerapi.JobSetTemplateSpec{
+				Spec: jobsetSpec,
+			},
+		},
+	}
+}
diff --git a/test/integration/singlecluster/controller/jobs/trainjob/trainjob_controller_test.go b/test/integration/singlecluster/controller/jobs/trainjob/trainjob_controller_test.go
index 36411c914..58c7bcfaf 100644
--- a/test/integration/singlecluster/controller/jobs/trainjob/trainjob_controller_test.go
+++ b/test/integration/singlecluster/controller/jobs/trainjob/trainjob_controller_test.go
@@ -312,16 +312,16 @@ var _ = ginkgo.Describe("TrainJob controller for workloads when only jobs with q
 				Replicas: 1,
 			}).
 			Obj()
-		testCtr := testingtrainjob.MakeClusterTrainingRuntime("test", testJobSet.Spec)
+		testTr := testingtrainjob.MakeTrainingRuntime("test", ns.Name, testJobSet.Spec)
 		trainJob := testingtrainjob.MakeTrainJob("trainjob-test", ns.Name).RuntimeRef(kftrainerapi.RuntimeRef{
 			APIGroup: ptr.To("trainer.kubeflow.org"),
 			Name:     "test",
-			Kind:     ptr.To("ClusterTrainingRuntime"),
+			Kind:     ptr.To(kftrainerapi.TrainingRuntimeKind),
 		}).
 			Suspend(false).
 			Obj()
 
-		util.MustCreate(ctx, k8sClient, testCtr)
+		util.MustCreate(ctx, k8sClient, testTr)
 		util.MustCreate(ctx, k8sClient, trainJob)
 		createdTrainJob := &kftrainerapi.TrainJob{}
 		gomega.Eventually(func(g gomega.Gomega) {
@@ -410,16 +410,16 @@ var _ = ginkgo.Describe("TrainJob controller interacting with scheduler", ginkgo
 			Request("node-1", corev1.ResourceCPU, "1").
 			Request("node-2", corev1.ResourceCPU, "1").
 			Obj()
-		testCtr := testingtrainjob.MakeClusterTrainingRuntime("test", testJobSet.Spec)
+		testTr := testingtrainjob.MakeTrainingRuntime("test", ns.Name, testJobSet.Spec)
 		trainJob := testingtrainjob.MakeTrainJob("trainjob-test", ns.Name).RuntimeRef(kftrainerapi.RuntimeRef{
 			APIGroup: ptr.To("trainer.kubeflow.org"),
 			Name:     "test",
-			Kind:     ptr.To("ClusterTrainingRuntime"),
+			Kind:     ptr.To(kftrainerapi.TrainingRuntimeKind),
 		}).
 			Queue("local-queue").
 			Obj()
 
-		util.MustCreate(ctx, k8sClient, testCtr)
+		util.MustCreate(ctx, k8sClient, testTr)
 		util.MustCreate(ctx, k8sClient, trainJob)
 		createdTrainJob := &kftrainerapi.TrainJob{}
 		gomega.Eventually(func(g gomega.Gomega) {
@@ -458,17 +458,17 @@ var _ = ginkgo.Describe("TrainJob controller interacting with scheduler", ginkgo
 			Request("node-1", corev1.ResourceCPU, "250m").
 			Request("node-2", corev1.ResourceCPU, "250m").
 			Obj()
-		testCtr1 := testingtrainjob.MakeClusterTrainingRuntime("ctr-1", testJobset1.Spec)
+		testTr1 := testingtrainjob.MakeTrainingRuntime("tr-1", ns.Name, testJobset1.Spec)
 		trainJob1 := testingtrainjob.MakeTrainJob("trainjob-test", ns.Name).RuntimeRef(kftrainerapi.RuntimeRef{
 			APIGroup: ptr.To("trainer.kubeflow.org"),
-			Name:     "ctr-1",
-			Kind:     ptr.To("ClusterTrainingRuntime"),
+			Name:     "tr-1",
+			Kind:     ptr.To(kftrainerapi.TrainingRuntimeKind),
 		}).
 			Queue("local-queue").
 			Suspend(true).
 			Obj()
 
-		util.MustCreate(ctx, k8sClient, testCtr1)
+		util.MustCreate(ctx, k8sClient, testTr1)
 		util.MustCreate(ctx, k8sClient, trainJob1)
 		ginkgo.By("checking the first trainjob starts", func() {
 			createdTrainJob1 := &kftrainerapi.TrainJob{}
@@ -497,17 +497,17 @@ var _ = ginkgo.Describe("TrainJob controller interacting with scheduler", ginkgo
 			Request("node-2", corev1.ResourceCPU, "1").
 			Obj()
 
-		testCtr2 := testingtrainjob.MakeClusterTrainingRuntime("ctr-2", testJobset2.Spec)
+		testTr2 := testingtrainjob.MakeTrainingRuntime("tr-2", ns.Name, testJobset2.Spec)
 		trainJob2 := testingtrainjob.MakeTrainJob("trainjob-test-2", ns.Name).RuntimeRef(kftrainerapi.RuntimeRef{
 			APIGroup: ptr.To("trainer.kubeflow.org"),
-			Name:     "ctr-2",
-			Kind:     ptr.To("ClusterTrainingRuntime"),
+			Name:     "tr-2",
+			Kind:     ptr.To(kftrainerapi.TrainingRuntimeKind),
 		}).
 			Queue("local-queue").
 			Suspend(true).
 			Obj()
 
-		util.MustCreate(ctx, k8sClient, testCtr2)
+		util.MustCreate(ctx, k8sClient, testTr2)
 		util.MustCreate(ctx, k8sClient, trainJob2)
 		ginkgo.By("checking a second no-fit trainjob does not start", func() {
 			createdTrainJob2 := &kftrainerapi.TrainJob{}
diff --git a/test/util/util.go b/test/util/util.go
index fb0f44757..ca821796c 100644
--- a/test/util/util.go
+++ b/test/util/util.go
@@ -131,6 +131,9 @@ func DeleteNamespace(ctx context.Context, c client.Client, ns *corev1.Namespace)
 	if err := DeleteAllJobsInNamespace(ctx, c, ns); err != nil {
 		return err
 	}
+	if err := DeleteAllTrainingRuntimesInNamespace(ctx, c, ns); err != nil {
+		return err
+	}
 	if err := c.DeleteAllOf(ctx, &appsv1.StatefulSet{}, client.InNamespace(ns.Name)); err != nil && !apierrors.IsNotFound(err) {
 		return err
 	}
@@ -169,6 +172,10 @@ func DeleteAllTrainJobsInNamespace(ctx context.Context, c client.Client, ns *cor
 	return deleteAllObjectsInNamespace(ctx, c, ns, &kftrainerapi.TrainJob{})
 }
 
+func DeleteAllTrainingRuntimesInNamespace(ctx context.Context, c client.Client, ns *corev1.Namespace) error {
+	return deleteAllObjectsInNamespace(ctx, c, ns, &kftrainerapi.TrainingRuntime{})
+}
+
 func DeleteAllLeaderWorkerSetsInNamespace(ctx context.Context, c client.Client, ns *corev1.Namespace) error {
 	return deleteAllObjectsInNamespace(ctx, c, ns, &leaderworkersetv1.LeaderWorkerSet{})
 }