Repository: kubernetes-sigs/kueue

Issue:
[flaky test] MultiKueue Should run a TrainJob on worker if admitted
/kind flake

**What happened**:

example: 
https://prow.k8s.io/view/gs/kubernetes-ci-logs/pr-logs/pull/kubernetes-sigs_kueue/7077/pull-kueue-test-integration-multikueue-main/1973001871303905280

**What you expected to happen**:
no failure
**How to reproduce it (as minimally and precisely as possible)**:
ci
**Anything else we need to know?**:

I think this started after re-enabling the test in https://github.com/kubernetes-sigs/kueue/pull/7068

unfortunately no more details:

```
/home/prow/go/src/sigs.k8s.io/kueue/test/integration/framework/framework.go:145
------------------------------
[DeferCleanup (Suite)] PASSED [0.001 seconds]
[DeferCleanup (Suite)] 
/home/prow/go/src/sigs.k8s.io/kueue/test/integration/framework/framework.go:145
------------------------------
[ReportAfterSuite] PASSED [2.849 seconds]
[ReportAfterSuite] Autogenerated ReportAfterSuite for --json-report --junit-report
autogenerated by Ginkgo
------------------------------
Ran 31 of 31 Specs in 169.067 seconds
SUCCESS! -- 31 Passed | 0 Failed | 0 Pending | 0 Skipped
Ginkgo ran 1 suite in 4m40.050654998s
Test Suite Failed
make: *** [Makefile-test.mk:90: test-multikueue-integration] Error 1

```

Discussion:
@mimowo: cc @ichekrygin @ichekrygin: Is this intermitted or more/less consistent repro?

Hoping, I can flush it out during https://github.com/kubernetes-sigs/kueue/pull/7095@mimowo: Unfortunately no easy repro for now. What is worse is that when it fails, then it does not say any specific reason.

We had a similar issue in the past where running with INTEGRATION_NPROC>1 would not output the DATA RACE of panics sometimes, so I opened https://github.com/kubernetes-sigs/kueue/pull/7085 but no luck so far@mimowo: @mszadkow I seem to remember you were investigating similarly failing multikueue integration tests (also without a clear indication of the reason). And for that reason we decreased INTEGRATION_NPROCS_MULTIKUEUE from 4 to 3. 

Do you remember maybe what was the underlying issue? 

I seem to remember this was after we added the Ray CRDs to the suite, maybe not this is happening after adding the Kubeflow Trainer v2 to the suite. An additional CRD may put the machines under stress more.

So, maybe this has nothing to do with the ElasticJobs which is just red herring here.@mimowo: Ok, I did a bit of archeology and for context we had this similar issue here: https://github.com/kubernetes-sigs/kueue/pull/3986, and as mitigation steps the proposal was summarized here: https://github.com/kubernetes-sigs/kueue/issues/4040:

```
For the experiments in the context of [Kuberay multikueue integrations tests](https://github.com/kubernetes-sigs/kueue/pull/3986) we learned that:
- increasing resources for the job for integration tests helped a little bit, but didn't fully solve the problem with running Ray (see [comment](https://github.com/kubernetes-sigs/kueue/pull/3986#issuecomment-2609155189))
- we know that setting INTEGRATION_NPROC=3 makes the tests usually pass, but it means lower parallelism enforced on the lightweight tests, which is not ideal
```
However, we have never fully understood why the issue was occurring, just found the above mitigates the issue.

I'm not yet totally sure what is the best way to move forward. I suppose the issue was triggered again by including the TrainerJob CRDs, but I would like to understand what is going on. Notably with decreased parallelism I wasn't able to repro here: https://github.com/kubernetes-sigs/kueue/pull/7085

If the issue continues and we cannot understand still maybe we need to mitigate again by decreasing parallelism (not sure yet, brainstorming).@mszadkow: I have been able to reproduce the issue, here is the log with race condition
`ginkgo --race -r --procs=4 --repeat=10`

```
All tests passed...
This was attempt 2 of 11.
Running Suite: MultiKueue Suite - /Users/michal_szadkowski/workspace/kueue/test/integration/multikueue
======================================================================================================
Random Seed: 1759314779

Will run 31 of 31 specs
Running in parallel across 4 processes
•••••••••••••••••••••••••••••
------------------------------
• [4.323 seconds]
MultiKueue Should run a TrainJob on worker if admitted
/Users/michal_szadkowski/workspace/kueue/test/integration/multikueue/jobs_test.go:1564

  Captured StdOut/StdErr Output >>
  ==================
  WARNING: DATA RACE
  Write at 0x000104cbc7a8 by goroutine 3233:
    github.com/kubeflow/trainer/v2/pkg/runtime/core.(*TrainingRuntime).newRuntimeInfo.WithPodSetSyncer.func6()
        /Users/michal_szadkowski/workspace/kueue/vendor/github.com/kubeflow/trainer/v2/pkg/runtime/runtime.go:173 +0x30
    github.com/kubeflow/trainer/v2/pkg/runtime.NewInfo()
        /Users/michal_szadkowski/workspace/kueue/vendor/github.com/kubeflow/trainer/v2/pkg/runtime/runtime.go:180 +0xf8
    github.com/kubeflow/trainer/v2/pkg/runtime/core.(*TrainingRuntime).newRuntimeInfo()
        /Users/michal_szadkowski/workspace/kueue/vendor/github.com/kubeflow/trainer/v2/pkg/runtime/core/trainingruntime.go:187 +0xe80
    github.com/kubeflow/trainer/v2/pkg/runtime/core.(*TrainingRuntime).RuntimeInfo()
        /Users/michal_szadkowski/workspace/kueue/vendor/github.com/kubeflow/trainer/v2/pkg/runtime/core/trainingruntime.go:103 +0xd8
    github.com/kubeflow/trainer/v2/pkg/runtime/core.(*ClusterTrainingRuntime).RuntimeInfo()
        /Users/michal_szadkowski/workspace/kueue/vendor/github.com/kubeflow/trainer/v2/pkg/runtime/core/clustertrainingruntime.go:71 +0x60
    sigs.k8s.io/kueue/pkg/controller/jobs/trainjob.getChildJobSet()
        /Users/michal_szadkowski/workspace/kueue/pkg/controller/jobs/trainjob/trainjob_controller.go:175 +0x1a8
    sigs.k8s.io/kueue/pkg/controller/jobs/trainjob.(*TrainJob).PodSets()
        /Users/michal_szadkowski/workspace/kueue/pkg/controller/jobs/trainjob/trainjob_controller.go:234 +0x28
    sigs.k8s.io/kueue/pkg/controller/jobframework.EquivalentToWorkload()
        /Users/michal_szadkowski/workspace/kueue/pkg/controller/jobframework/reconciler.go:1047 +0x144
    sigs.k8s.io/kueue/pkg/controller/jobframework.FindMatchingWorkloads()
        /Users/michal_szadkowski/workspace/kueue/pkg/controller/jobframework/reconciler.go:940 +0x308
    sigs.k8s.io/kueue/pkg/controller/jobframework.(*JobReconciler).ensureOneWorkload()
        /Users/michal_szadkowski/workspace/kueue/pkg/controller/jobframework/reconciler.go:859 +0x6c8
    sigs.k8s.io/kueue/pkg/controller/jobframework.(*JobReconciler).ReconcileGenericJob()
        /Users/michal_szadkowski/workspace/kueue/pkg/controller/jobframework/reconciler.go:397 +0xf2c
    sigs.k8s.io/kueue/pkg/controller/jobs/trainjob.(*trainJobReconciler).Reconcile()
        /Users/michal_szadkowski/workspace/kueue/pkg/controller/jobs/trainjob/trainjob_controller.go:106 +0xd0
    sigs.k8s.io/controller-runtime/pkg/internal/controller.(*Controller[go.shape.struct { k8s.io/apimachinery/pkg/types.NamespacedName }]).Reconcile()
        /Users/michal_szadkowski/workspace/kueue/vendor/sigs.k8s.io/controller-runtime/pkg/internal/controller/controller.go:216 +0x164
    sigs.k8s.io/controller-runtime/pkg/internal/controller.(*Controller[go.shape.struct { k8s.io/apimachinery/pkg/types.NamespacedName }]).reconcileHandler()
        /Users/michal_szadkowski/workspace/kueue/vendor/sigs.k8s.io/controller-runtime/pkg/internal/controller/controller.go:461 +0x374
    sigs.k8s.io/controller-runtime/pkg/internal/controller.(*Controller[go.shape.struct { k8s.io/apimachinery/pkg/types.NamespacedName }]).processNextWorkItem()
        /Users/michal_szadkowski/workspace/kueue/vendor/sigs.k8s.io/controller-runtime/pkg/internal/controller/controller.go:421 +0x23c
    sigs.k8s.io/controller-runtime/pkg/internal/controller.(*Controller[go.shape.struct { k8s.io/apimachinery/pkg/types.NamespacedName }]).Start.func1.1()
        /Users/michal_szadkowski/workspace/kueue/vendor/sigs.k8s.io/controller-runtime/pkg/internal/controller/controller.go:296 +0xb8

  Previous write at 0x000104cbc7a8 by goroutine 3133:
    github.com/kubeflow/trainer/v2/pkg/runtime/core.(*TrainingRuntime).newRuntimeInfo.WithPodSetSyncer.func6()
        /Users/michal_szadkowski/workspace/kueue/vendor/github.com/kubeflow/trainer/v2/pkg/runtime/runtime.go:173 +0x30
    github.com/kubeflow/trainer/v2/pkg/runtime.NewInfo()
        /Users/michal_szadkowski/workspace/kueue/vendor/github.com/kubeflow/trainer/v2/pkg/runtime/runtime.go:180 +0xf8
    github.com/kubeflow/trainer/v2/pkg/runtime/core.(*TrainingRuntime).newRuntimeInfo()
        /Users/michal_szadkowski/workspace/kueue/vendor/github.com/kubeflow/trainer/v2/pkg/runtime/core/trainingruntime.go:187 +0xe80
    github.com/kubeflow/trainer/v2/pkg/runtime/core.(*TrainingRuntime).RuntimeInfo()
        /Users/michal_szadkowski/workspace/kueue/vendor/github.com/kubeflow/trainer/v2/pkg/runtime/core/trainingruntime.go:103 +0xd8
    github.com/kubeflow/trainer/v2/pkg/runtime/core.(*ClusterTrainingRuntime).RuntimeInfo()
        /Users/michal_szadkowski/workspace/kueue/vendor/github.com/kubeflow/trainer/v2/pkg/runtime/core/clustertrainingruntime.go:71 +0x60
    sigs.k8s.io/kueue/pkg/controller/jobs/trainjob.getChildJobSet()
        /Users/michal_szadkowski/workspace/kueue/pkg/controller/jobs/trainjob/trainjob_controller.go:175 +0x1a8
    sigs.k8s.io/kueue/pkg/controller/jobs/trainjob.(*TrainJob).PodSets()
        /Users/michal_szadkowski/workspace/kueue/pkg/controller/jobs/trainjob/trainjob_controller.go:234 +0x28
    sigs.k8s.io/kueue/pkg/controller/jobframework.EquivalentToWorkload()
        /Users/michal_szadkowski/workspace/kueue/pkg/controller/jobframework/reconciler.go:1047 +0x144
    sigs.k8s.io/kueue/pkg/controller/jobframework.FindMatchingWorkloads()
        /Users/michal_szadkowski/workspace/kueue/pkg/controller/jobframework/reconciler.go:940 +0x308
    sigs.k8s.io/kueue/pkg/controller/jobframework.(*JobReconciler).ensureOneWorkload()
        /Users/michal_szadkowski/workspace/kueue/pkg/controller/jobframework/reconciler.go:859 +0x6c8
    sigs.k8s.io/kueue/pkg/controller/jobframework.(*JobReconciler).ReconcileGenericJob()
        /Users/michal_szadkowski/workspace/kueue/pkg/controller/jobframework/reconciler.go:397 +0xf2c
    sigs.k8s.io/kueue/pkg/controller/jobs/trainjob.(*trainJobReconciler).Reconcile()
        /Users/michal_szadkowski/workspace/kueue/pkg/controller/jobs/trainjob/trainjob_controller.go:106 +0xd0
    sigs.k8s.io/controller-runtime/pkg/internal/controller.(*Controller[go.shape.struct { k8s.io/apimachinery/pkg/types.NamespacedName }]).Reconcile()
        /Users/michal_szadkowski/workspace/kueue/vendor/sigs.k8s.io/controller-runtime/pkg/internal/controller/controller.go:216 +0x164
    sigs.k8s.io/controller-runtime/pkg/internal/controller.(*Controller[go.shape.struct { k8s.io/apimachinery/pkg/types.NamespacedName }]).reconcileHandler()
        /Users/michal_szadkowski/workspace/kueue/vendor/sigs.k8s.io/controller-runtime/pkg/internal/controller/controller.go:461 +0x374
    sigs.k8s.io/controller-runtime/pkg/internal/controller.(*Controller[go.shape.struct { k8s.io/apimachinery/pkg/types.NamespacedName }]).processNextWorkItem()
        /Users/michal_szadkowski/workspace/kueue/vendor/sigs.k8s.io/controller-runtime/pkg/internal/controller/controller.go:421 +0x23c
    sigs.k8s.io/controller-runtime/pkg/internal/controller.(*Controller[go.shape.struct { k8s.io/apimachinery/pkg/types.NamespacedName }]).Start.func1.1()
        /Users/michal_szadkowski/workspace/kueue/vendor/sigs.k8s.io/controller-runtime/pkg/internal/controller/controller.go:296 +0xb8

  Goroutine 3233 (running) created at:
    sigs.k8s.io/controller-runtime/pkg/internal/controller.(*Controller[go.shape.struct { k8s.io/apimachinery/pkg/types.NamespacedName }]).Start.func1()
        /Users/michal_szadkowski/workspace/kueue/vendor/sigs.k8s.io/controller-runtime/pkg/internal/controller/controller.go:292 +0x2d8
    sigs.k8s.io/controller-runtime/pkg/internal/controller.(*Controller[go.shape.struct { k8s.io/apimachinery/pkg/types.NamespacedName }]).Start()
        /Users/michal_szadkowski/workspace/kueue/vendor/sigs.k8s.io/controller-runtime/pkg/internal/controller/controller.go:303 +0x14c
    sigs.k8s.io/controller-runtime/pkg/internal/controller.(*Controller[sigs.k8s.io/controller-runtime/pkg/reconcile.Request]).Start()
        /Users/michal_szadkowski/workspace/kueue/vendor/sigs.k8s.io/controller-runtime/pkg/internal/controller/controller.go:259 +0x48
    sigs.k8s.io/controller-runtime/pkg/manager.(*runnableGroup).reconcile.func1()
        /Users/michal_szadkowski/workspace/kueue/vendor/sigs.k8s.io/controller-runtime/pkg/manager/runnable_group.go:260 +0x158
    sigs.k8s.io/controller-runtime/pkg/manager.(*runnableGroup).reconcile.gowrap1()
        /Users/michal_szadkowski/workspace/kueue/vendor/sigs.k8s.io/controller-runtime/pkg/manager/runnable_group.go:283 +0x40

  Goroutine 3133 (running) created at:
    sigs.k8s.io/controller-runtime/pkg/internal/controller.(*Controller[go.shape.struct { k8s.io/apimachinery/pkg/types.NamespacedName }]).Start.func1()
        /Users/michal_szadkowski/workspace/kueue/vendor/sigs.k8s.io/controller-runtime/pkg/internal/controller/controller.go:292 +0x2d8
    sigs.k8s.io/controller-runtime/pkg/internal/controller.(*Controller[go.shape.struct { k8s.io/apimachinery/pkg/types.NamespacedName }]).Start()
        /Users/michal_szadkowski/workspace/kueue/vendor/sigs.k8s.io/controller-runtime/pkg/internal/controller/controller.go:303 +0x14c
    sigs.k8s.io/controller-runtime/pkg/internal/controller.(*Controller[sigs.k8s.io/controller-runtime/pkg/reconcile.Request]).Start()
        /Users/michal_szadkowski/workspace/kueue/vendor/sigs.k8s.io/controller-runtime/pkg/internal/controller/controller.go:259 +0x48
    sigs.k8s.io/controller-runtime/pkg/manager.(*runnableGroup).reconcile.func1()
        /Users/michal_szadkowski/workspace/kueue/vendor/sigs.k8s.io/controller-runtime/pkg/manager/runnable_group.go:260 +0x158
    sigs.k8s.io/controller-runtime/pkg/manager.(*runnableGroup).reconcile.gowrap1()
        /Users/michal_szadkowski/workspace/kueue/vendor/sigs.k8s.io/controller-runtime/pkg/manager/runnable_group.go:283 +0x40
  ==================
  << Captured StdOut/StdErr Output
------------------------------
•

Ran 31 of 31 Specs in 147.607 seconds
SUCCESS! -- 31 Passed | 0 Failed | 0 Pending | 0 Skipped


Tests failed on attempt #3


Ginkgo ran 1 suite in 7m49.903828292s

Test Suite Failed
(prow-artifacts) [12:35:27 @michal_szadkowski:~/workspace/kueue/test/integration/multikueue]
```@mszadkow: Well this is unfortunate that race conditions are not visible when we run tests with Makefile.
However I do believe this is different from https://github.com/kubernetes-sigs/kueue/issues/3987 
There we had a panic and test suite stopped working all together.@mimowo: nice, thank you for the repro. I see the TrainJob on the stracktrace which is a new card indeed. Do you know which specific It (test case) triggers the race?

cc @kaisoz @mszadkow: I think this is this one:
```
MultiKueue Should run a TrainJob on worker if admitted
kueue/test/integration/multikueue/jobs_test.go:1564
```@mimowo: I see, so this looks like two threads are modifying the internal state of the `runtime` somewhere here: https://github.com/kubernetes-sigs/kueue/blob/81113ca514a6b20dce3e0b3fecc6d6fe14b75109/pkg/controller/jobs/trainjob/trainjob_controller.go#L176C15-L176C43

I guess we may need to DeepCopy the runtime before calling RuntimeInfo. Or avoid storing "runtimes" as state, seems like premature optimization (potentially). If you have a reliable repro then worth checking if DeepCopy or getting new runtime each time helps.

EDIT: if that does not work we may need to take a mutex on the getChildJobSet function (or even just around calling RuntimeInfo).@mszadkow: /assign@mimowo: /retitle [flaky test] MultiKueue Should run a TrainJob on worker if admitted@mimowo: cc @tenzen-y @andreyvelich I'm also wondering if TrainerJob could / should internally avoid mutating the state by calling RuntimeInfo.@mszadkow: Actually if we talk about the first idea, what it could be, I learnt that is quite easy to reproduce, by increasing number of processes
When we run: `ginkgo --race -r --procs=6 --repeat=10`
We got old issue with CRD everytime:
```
  [FAILED] in [BeforeSuite] - /Users/michal_szadkowski/workspace/kueue/vendor/github.com/onsi/ginkgo/v2/internal/suite.go:328 @ 10/01/25 15:14:37.505
  << Timeline

  [FAILED] Unexpected error:
      <*fmt.wrapError | 0xc004dcfbc0>: 
      unable to install CRDs onto control plane: something went wrong waiting for CRDs to appear as API resources: context deadline exceeded
      {
          msg: "unable to install CRDs onto control plane: something went wrong waiting for CRDs to appear as API resources: context deadline exceeded",
          err: <*fmt.wrapError | 0xc004dcfb80>{
              msg: "something went wrong waiting for CRDs to appear as API resources: context deadline exceeded",
              err: <context.deadlineExceededError>{},
          },
      }
  occurred
  In [BeforeSuite] at: /Users/michal_szadkowski/workspace/kueue/vendor/github.com/onsi/ginkgo/v2/internal/suite.go:328 @ 10/01/25 15:14:37.505
------------------------------
[AfterSuite] [PANICKED] [0.006 seconds]
[AfterSuite] 
/Users/michal_szadkowski/workspace/kueue/test/integration/multikueue/suite_test.go:396

  Timeline >>
  STEP: tearing down the test environment @ 10/01/25 15:14:37.816
  [PANICKED] in [AfterSuite] - /Users/michal_szadkowski/go/pkg/mod/golang.org/toolchain@v0.0.1-go1.25.0.darwin-arm64/src/runtime/panic.go:262 @ 10/01/25 15:14:37.817
  << Timeline

  [PANICKED] Test Panicked
  In [AfterSuite] at: /Users/michal_szadkowski/go/pkg/mod/golang.org/toolchain@v0.0.1-go1.25.0.darwin-arm64/src/runtime/panic.go:262 @ 10/01/25 15:14:37.817

  runtime error: invalid memory address or nil pointer dereference

  Full Stack Trace
    sigs.k8s.io/kueue/test/integration/framework.(*Framework).Teardown(0x0)
        /Users/michal_szadkowski/workspace/kueue/test/integration/framework/framework.go:218 +0x44
    sigs.k8s.io/kueue/test/integration/multikueue.init.func10()
        /Users/michal_szadkowski/workspace/kueue/test/integration/multikueue/suite_test.go:397 +0x38
------------------------------
•••••••••••••••••••••••••••••••

Summarizing 2 Failures:
  [FAIL] [BeforeSuite] 
  /Users/michal_szadkowski/workspace/kueue/vendor/github.com/onsi/ginkgo/v2/internal/suite.go:328
  [PANICKED!] [AfterSuite] 
  /Users/michal_szadkowski/go/pkg/mod/golang.org/toolchain@v0.0.1-go1.25.0.darwin-arm64/src/runtime/panic.go:262

Ran 31 of 31 Specs in 170.093 seconds
FAIL! -- 31 Passed | 0 Failed | 0 Pending | 0 Skipped
```
@mimowo: nice! @kaisoz: > ...Or avoid storing "runtimes" as state, seems like premature optimization (potentially). If you have a reliable repro then worth checking if DeepCopy or getting new runtime each time helps.

Unfortunately the runtimes can ony be created once at this moment (I've tried it already)

The call to the runtimes [New](https://github.com/kubernetes-sigs/kueue/blob/81113ca514a6b20dce3e0b3fecc6d6fe14b75109/pkg/controller/jobs/trainjob/trainjob_controller.go#L93) function  (implementation [here](https://github.com/kubeflow/trainer/blob/b918411a6b4d5dfc4e68a328139cdb7d821040f2/pkg/runtime/core/core.go#L31)) ends up calling each runtime constructor, which in turn [register the field in the indexer](https://github.com/kubeflow/trainer/blob/b918411a6b4d5dfc4e68a328139cdb7d821040f2/pkg/runtime/core/trainingruntime.go#L65). This registration **can only be done once**

@mszadkow: > > ...Or avoid storing "runtimes" as state, seems like premature optimization (potentially). If you have a reliable repro then worth checking if DeepCopy or getting new runtime each time helps.
> 
> Unfortunately the runtimes can ony be created once at this moment (I've tried it already)
> 
> The call to the runtimes [New](https://github.com/kubernetes-sigs/kueue/blob/81113ca514a6b20dce3e0b3fecc6d6fe14b75109/pkg/controller/jobs/trainjob/trainjob_controller.go#L93) function (implementation [here](https://github.com/kubeflow/trainer/blob/b918411a6b4d5dfc4e68a328139cdb7d821040f2/pkg/runtime/core/core.go#L31)) ends up calling each runtime constructor, which in turn [register the field in the indexer](https://github.com/kubeflow/trainer/blob/b918411a6b4d5dfc4e68a328139cdb7d821040f2/pkg/runtime/core/trainingruntime.go#L65). This registration **can only be done once**

I guess lock on `trainjob.getChildJobSet` is the best option here@mimowo: > I guess lock on trainjob.getChildJobSet is the best option here

I think so, at least as a quick fix. We can report in the Kubeflow Trainer to make it thread safe. 
Also, I think you don't need to get mutex for entire getChildJobSet, just calling RuntimeInfo (could be extracted to another helper wrapped in mutex like), runtimeInfo.

Solution patch:
diff --git a/Dockerfile b/Dockerfile
index 36eb3ce44..4a865b09a 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,5 +1,5 @@
 # defaulted ARGs need to be declared first
-ARG BUILDER_IMAGE=golang:1.24
+ARG BUILDER_IMAGE=golang:1.25
 ARG BASE_IMAGE=gcr.io/distroless/static:nonroot
 # compilation stage for the manager binary
 FROM --platform=${BUILDPLATFORM} ${BUILDER_IMAGE} AS builder
diff --git a/charts/kueue/templates/webhook/manifests.yaml b/charts/kueue/templates/webhook/manifests.yaml
index d2ae17a24..784e03f67 100644
--- a/charts/kueue/templates/webhook/manifests.yaml
+++ b/charts/kueue/templates/webhook/manifests.yaml
@@ -242,26 +242,6 @@ webhooks:
           - xgboostjobs
     sideEffects: None
     reinvocationPolicy: '{{ .Values.mutatingWebhook.reinvocationPolicy }}'
-  - admissionReviewVersions:
-      - v1
-    clientConfig:
-      service:
-        name: '{{ include "kueue.fullname" . }}-webhook-service'
-        namespace: '{{ .Release.Namespace }}'
-        path: /mutate-trainer-kubeflow-org-v1alpha1-trainjob
-    failurePolicy: Fail
-    name: mtrainjob.kb.io
-    rules:
-      - apiGroups:
-          - trainer.kubeflow.org
-        apiVersions:
-          - v1alpha1
-        operations:
-          - CREATE
-        resources:
-          - trainjobs
-    sideEffects: None
-    reinvocationPolicy: '{{ .Values.mutatingWebhook.reinvocationPolicy }}'
   - admissionReviewVersions:
       - v1
     clientConfig:
@@ -430,6 +410,26 @@ webhooks:
           - statefulsets
     sideEffects: None
     reinvocationPolicy: '{{ .Values.mutatingWebhook.reinvocationPolicy }}'
+  - admissionReviewVersions:
+      - v1
+    clientConfig:
+      service:
+        name: '{{ include "kueue.fullname" . }}-webhook-service'
+        namespace: '{{ .Release.Namespace }}'
+        path: /mutate-trainer-kubeflow-org-v1alpha1-trainjob
+    failurePolicy: Fail
+    name: mtrainjob.kb.io
+    rules:
+      - apiGroups:
+          - trainer.kubeflow.org
+        apiVersions:
+          - v1alpha1
+        operations:
+          - CREATE
+        resources:
+          - trainjobs
+    sideEffects: None
+    reinvocationPolicy: '{{ .Values.mutatingWebhook.reinvocationPolicy }}'
   - admissionReviewVersions:
       - v1
     clientConfig:
@@ -730,26 +730,6 @@ webhooks:
         resources:
           - xgboostjobs
     sideEffects: None
-  - admissionReviewVersions:
-      - v1
-    clientConfig:
-      service:
-        name: '{{ include "kueue.fullname" . }}-webhook-service'
-        namespace: '{{ .Release.Namespace }}'
-        path: /validate-trainer-kubeflow-org-v1alpha1-trainjob
-    failurePolicy: Fail
-    name: vtrainjob.kb.io
-    rules:
-      - apiGroups:
-          - trainer.kubeflow.org
-        apiVersions:
-          - v1alpha1
-        operations:
-          - CREATE
-          - UPDATE
-        resources:
-          - trainjobs
-    sideEffects: None
   - admissionReviewVersions:
       - v1
     clientConfig:
@@ -916,6 +896,26 @@ webhooks:
         resources:
           - statefulsets
     sideEffects: None
+  - admissionReviewVersions:
+      - v1
+    clientConfig:
+      service:
+        name: '{{ include "kueue.fullname" . }}-webhook-service'
+        namespace: '{{ .Release.Namespace }}'
+        path: /validate-trainer-kubeflow-org-v1alpha1-trainjob
+    failurePolicy: Fail
+    name: vtrainjob.kb.io
+    rules:
+      - apiGroups:
+          - trainer.kubeflow.org
+        apiVersions:
+          - v1alpha1
+        operations:
+          - CREATE
+          - UPDATE
+        resources:
+          - trainjobs
+    sideEffects: None
   - admissionReviewVersions:
       - v1
     clientConfig:
diff --git a/cmd/kueueviz/backend/go.mod b/cmd/kueueviz/backend/go.mod
index fc7b3fcb0..81dcc25f0 100644
--- a/cmd/kueueviz/backend/go.mod
+++ b/cmd/kueueviz/backend/go.mod
@@ -1,6 +1,6 @@
 module kueueviz
 
-go 1.24.0
+go 1.25.0
 
 require (
 	github.com/gin-contrib/cors v1.7.6
diff --git a/config/components/webhook/manifests.yaml b/config/components/webhook/manifests.yaml
index ca27ba224..26caf441f 100644
--- a/config/components/webhook/manifests.yaml
+++ b/config/components/webhook/manifests.yaml
@@ -176,25 +176,6 @@ webhooks:
     resources:
     - xgboostjobs
   sideEffects: None
-- admissionReviewVersions:
-  - v1
-  clientConfig:
-    service:
-      name: webhook-service
-      namespace: system
-      path: /mutate-trainer-kubeflow-org-v1alpha1-trainjob
-  failurePolicy: Fail
-  name: mtrainjob.kb.io
-  rules:
-  - apiGroups:
-    - trainer.kubeflow.org
-    apiVersions:
-    - v1alpha1
-    operations:
-    - CREATE
-    resources:
-    - trainjobs
-  sideEffects: None
 - admissionReviewVersions:
   - v1
   clientConfig:
@@ -311,6 +292,25 @@ webhooks:
     resources:
     - statefulsets
   sideEffects: None
+- admissionReviewVersions:
+  - v1
+  clientConfig:
+    service:
+      name: webhook-service
+      namespace: system
+      path: /mutate-trainer-kubeflow-org-v1alpha1-trainjob
+  failurePolicy: Fail
+  name: mtrainjob.kb.io
+  rules:
+  - apiGroups:
+    - trainer.kubeflow.org
+    apiVersions:
+    - v1alpha1
+    operations:
+    - CREATE
+    resources:
+    - trainjobs
+  sideEffects: None
 - admissionReviewVersions:
   - v1
   clientConfig:
@@ -554,26 +554,6 @@ webhooks:
     resources:
     - xgboostjobs
   sideEffects: None
-- admissionReviewVersions:
-  - v1
-  clientConfig:
-    service:
-      name: webhook-service
-      namespace: system
-      path: /validate-trainer-kubeflow-org-v1alpha1-trainjob
-  failurePolicy: Fail
-  name: vtrainjob.kb.io
-  rules:
-  - apiGroups:
-    - trainer.kubeflow.org
-    apiVersions:
-    - v1alpha1
-    operations:
-    - CREATE
-    - UPDATE
-    resources:
-    - trainjobs
-  sideEffects: None
 - admissionReviewVersions:
   - v1
   clientConfig:
@@ -694,6 +674,26 @@ webhooks:
     resources:
     - statefulsets
   sideEffects: None
+- admissionReviewVersions:
+  - v1
+  clientConfig:
+    service:
+      name: webhook-service
+      namespace: system
+      path: /validate-trainer-kubeflow-org-v1alpha1-trainjob
+  failurePolicy: Fail
+  name: vtrainjob.kb.io
+  rules:
+  - apiGroups:
+    - trainer.kubeflow.org
+    apiVersions:
+    - v1alpha1
+    operations:
+    - CREATE
+    - UPDATE
+    resources:
+    - trainjobs
+  sideEffects: None
 - admissionReviewVersions:
   - v1
   clientConfig:
diff --git a/go.mod b/go.mod
index 37a3cc9e3..04d809c95 100644
--- a/go.mod
+++ b/go.mod
@@ -1,6 +1,6 @@
 module sigs.k8s.io/kueue
 
-go 1.24.0
+go 1.25.0
 
 require (
 	github.com/cert-manager/cert-manager v1.18.2
diff --git a/hack/internal/tools/go.mod b/hack/internal/tools/go.mod
index eae968f48..bc32eb66f 100644
--- a/hack/internal/tools/go.mod
+++ b/hack/internal/tools/go.mod
@@ -1,6 +1,6 @@
 module sigs.k8s.io/kueue/internal/tools
 
-go 1.24.0
+go 1.25.0
 
 tool (
 	github.com/gohugoio/hugo
diff --git a/pkg/controller/jobs/jobs.go b/pkg/controller/jobs/jobs.go
index f9c38fca5..e3ca17b66 100644
--- a/pkg/controller/jobs/jobs.go
+++ b/pkg/controller/jobs/jobs.go
@@ -23,11 +23,11 @@ import (
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/job"
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/jobset"
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/jobs"
-	_ "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/trainjob"
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/leaderworkerset"
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/mpijob"
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/pod"
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/raycluster"
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/rayjob"
 	_ "sigs.k8s.io/kueue/pkg/controller/jobs/statefulset"
+	_ "sigs.k8s.io/kueue/pkg/controller/jobs/trainjob"
 )
diff --git a/pkg/controller/jobs/kubeflow/trainjob/trainjob_controller.go b/pkg/controller/jobs/trainjob/trainjob_controller.go
similarity index 90%
rename from pkg/controller/jobs/kubeflow/trainjob/trainjob_controller.go
rename to pkg/controller/jobs/trainjob/trainjob_controller.go
index 87ee6438c..0d36ec091 100644
--- a/pkg/controller/jobs/kubeflow/trainjob/trainjob_controller.go
+++ b/pkg/controller/jobs/trainjob/trainjob_controller.go
@@ -23,7 +23,7 @@ import (
 	"fmt"
 	"strconv"
 
-	kftrainerapi "github.com/kubeflow/trainer/v2/pkg/apis/trainer/v1alpha1"
+	kftrainer "github.com/kubeflow/trainer/v2/pkg/apis/trainer/v1alpha1"
 	kftrainerruntime "github.com/kubeflow/trainer/v2/pkg/runtime"
 	kftrainerruntimecore "github.com/kubeflow/trainer/v2/pkg/runtime/core"
 	kftrainerjobset "github.com/kubeflow/trainer/v2/pkg/runtime/framework/plugins/jobset"
@@ -42,7 +42,7 @@ import (
 
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
-	kJobset "sigs.k8s.io/kueue/pkg/controller/jobs/jobset"
+	workloadjobset "sigs.k8s.io/kueue/pkg/controller/jobs/jobset"
 	"sigs.k8s.io/kueue/pkg/features"
 	"sigs.k8s.io/kueue/pkg/podset"
 	clientutil "sigs.k8s.io/kueue/pkg/util/client"
@@ -50,7 +50,7 @@ import (
 )
 
 var (
-	gvk                    = kftrainerapi.GroupVersion.WithKind("TrainJob")
+	gvk                    = kftrainer.GroupVersion.WithKind("TrainJob")
 	FrameworkName          = "trainer.kubeflow.org/trainjob"
 	TrainJobControllerName = "trainer.kubeflow.org/trainjob-controller"
 )
@@ -66,8 +66,8 @@ func init() {
 		NewJob:            NewJob,
 		NewReconciler:     NewReconciler,
 		SetupWebhook:      SetupTrainJobWebhook,
-		JobType:           &kftrainerapi.TrainJob{},
-		AddToScheme:       kftrainerapi.AddToScheme,
+		JobType:           &kftrainer.TrainJob{},
+		AddToScheme:       kftrainer.AddToScheme,
 		MultiKueueAdapter: &multiKueueAdapter{},
 	}))
 }
@@ -108,11 +108,11 @@ func (r *trainJobReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
 
 func (r *trainJobReconciler) SetupWithManager(mgr ctrl.Manager) error {
 	b := ctrl.NewControllerManagedBy(mgr).
-		For(&kftrainerapi.TrainJob{}).Owns(&kueue.Workload{}).Owns(&jobsetapi.JobSet{})
+		For(&kftrainer.TrainJob{}).Owns(&kueue.Workload{}).Owns(&jobsetapi.JobSet{})
 	return b.Complete(r)
 }
 
-type TrainJob kftrainerapi.TrainJob
+type TrainJob kftrainer.TrainJob
 
 var _ jobframework.GenericJob = (*TrainJob)(nil)
 var _ jobframework.JobWithCustomStop = (*TrainJob)(nil)
@@ -124,11 +124,11 @@ func NewJob() jobframework.GenericJob {
 }
 
 func fromObject(obj runtime.Object) *TrainJob {
-	return (*TrainJob)(obj.(*kftrainerapi.TrainJob))
+	return (*TrainJob)(obj.(*kftrainer.TrainJob))
 }
 
 func (t *TrainJob) Object() client.Object {
-	return (*kftrainerapi.TrainJob)(t)
+	return (*kftrainer.TrainJob)(t)
 }
 
 func (t *TrainJob) IsSuspended() bool {
@@ -167,7 +167,7 @@ func getChildJobSet(t *TrainJob) (*jobsetapi.JobSet, error) {
 		return nil, fmt.Errorf("unsupported runtime: %s", runtimeRefGK)
 	}
 
-	trainJob := (*kftrainerapi.TrainJob)(t)
+	trainJob := (*kftrainer.TrainJob)(t)
 	trSpec, err := getRuntimeSpec(trainJob)
 	if err != nil {
 		return nil, fmt.Errorf("runtime '%s' not found", trainJob.Spec.RuntimeRef.Name)
@@ -212,16 +212,16 @@ func jobsetApplyToJobset(jobsetApply *jobsetapplyapi.JobSetApplyConfiguration) (
 	return jobset, nil
 }
 
-func getRuntimeSpec(trainJob *kftrainerapi.TrainJob) (*kftrainerapi.TrainingRuntimeSpec, error) {
-	if *trainJob.Spec.RuntimeRef.Kind == kftrainerapi.ClusterTrainingRuntimeKind {
-		var ctr kftrainerapi.ClusterTrainingRuntime
+func getRuntimeSpec(trainJob *kftrainer.TrainJob) (*kftrainer.TrainingRuntimeSpec, error) {
+	if *trainJob.Spec.RuntimeRef.Kind == kftrainer.ClusterTrainingRuntimeKind {
+		var ctr kftrainer.ClusterTrainingRuntime
 		err := reconciler.client.Get(reconciler.ctx, client.ObjectKey{Name: trainJob.Spec.RuntimeRef.Name}, &ctr)
 		if err != nil {
 			return nil, err
 		}
 		return &ctr.Spec, nil
 	} else {
-		var tr kftrainerapi.TrainingRuntime
+		var tr kftrainer.TrainingRuntime
 		err := reconciler.client.Get(reconciler.ctx, client.ObjectKey{Namespace: trainJob.Namespace, Name: trainJob.Spec.RuntimeRef.Name}, &tr)
 		if err != nil {
 			return nil, err
@@ -235,7 +235,7 @@ func (t *TrainJob) PodSets() ([]kueue.PodSet, error) {
 	if err != nil {
 		return nil, err
 	}
-	return (*kJobset.JobSet)(jobset).PodSets()
+	return (*workloadjobset.JobSet)(jobset).PodSets()
 }
 
 func (t *TrainJob) RunWithPodSetsInfo(podSetsInfo []podset.PodSetInfo) error {
@@ -249,7 +249,7 @@ func (t *TrainJob) RunWithPodSetsInfo(podSetsInfo []podset.PodSetInfo) error {
 	}
 
 	if t.Spec.PodSpecOverrides == nil {
-		t.Spec.PodSpecOverrides = []kftrainerapi.PodSpecOverride{}
+		t.Spec.PodSpecOverrides = []kftrainer.PodSpecOverride{}
 	}
 	if t.Annotations == nil {
 		t.Annotations = map[string]string{}
@@ -257,8 +257,8 @@ func (t *TrainJob) RunWithPodSetsInfo(podSetsInfo []podset.PodSetInfo) error {
 	t.Annotations[firstOverrideIdx] = strconv.Itoa(len(t.Spec.PodSpecOverrides))
 	for _, info := range podSetsInfo {
 		// The trainjob controller merges each podSpecOverride sequentially, so any existing user provided override will be processed first
-		t.Spec.PodSpecOverrides = append(t.Spec.PodSpecOverrides, kftrainerapi.PodSpecOverride{
-			TargetJobs: []kftrainerapi.PodSpecOverrideTargetJob{
+		t.Spec.PodSpecOverrides = append(t.Spec.PodSpecOverrides, kftrainer.PodSpecOverride{
+			TargetJobs: []kftrainer.PodSpecOverrideTargetJob{
 				{Name: string(info.Name)},
 			},
 			// TODO: Set the labels/annotations when supported. See https://github.com/kubeflow/trainer/pull/2785
@@ -320,10 +320,10 @@ func (t *TrainJob) RestorePodSetsInfo(_ []podset.PodSetInfo) bool {
 }
 
 func (t *TrainJob) Finished() (message string, success, finished bool) {
-	if c := apimeta.FindStatusCondition(t.Status.Conditions, kftrainerapi.TrainJobComplete); c != nil && c.Status == metav1.ConditionTrue {
+	if c := apimeta.FindStatusCondition(t.Status.Conditions, kftrainer.TrainJobComplete); c != nil && c.Status == metav1.ConditionTrue {
 		return c.Message, true, true
 	}
-	if c := apimeta.FindStatusCondition(t.Status.Conditions, kftrainerapi.TrainJobFailed); c != nil && c.Status == metav1.ConditionTrue {
+	if c := apimeta.FindStatusCondition(t.Status.Conditions, kftrainer.TrainJobFailed); c != nil && c.Status == metav1.ConditionTrue {
 		return c.Message, false, true
 	}
 	return message, success, false
@@ -356,7 +356,7 @@ func (t *TrainJob) ReclaimablePods() ([]kueue.ReclaimablePod, error) {
 	}
 
 	ret := make([]kueue.ReclaimablePod, 0, len(jobset.Spec.ReplicatedJobs))
-	statuses := slices.ToRefMap(t.Status.JobsStatus, func(js *kftrainerapi.JobStatus) string { return js.Name })
+	statuses := slices.ToRefMap(t.Status.JobsStatus, func(js *kftrainer.JobStatus) string { return js.Name })
 
 	for i := range jobset.Spec.ReplicatedJobs {
 		spec := &jobset.Spec.ReplicatedJobs[i]
@@ -364,7 +364,7 @@ func (t *TrainJob) ReclaimablePods() ([]kueue.ReclaimablePod, error) {
 			if status.Succeeded > 0 && status.Succeeded <= spec.Replicas {
 				ret = append(ret, kueue.ReclaimablePod{
 					Name:  kueue.NewPodSetReference(spec.Name),
-					Count: status.Succeeded * kJobset.PodsCountPerReplica(spec),
+					Count: status.Succeeded * workloadjobset.PodsCountPerReplica(spec),
 				})
 			}
 		}
diff --git a/pkg/controller/jobs/kubeflow/trainjob/trainjob_multikueue_adapter.go b/pkg/controller/jobs/trainjob/trainjob_multikueue_adapter.go
similarity index 100%
rename from pkg/controller/jobs/kubeflow/trainjob/trainjob_multikueue_adapter.go
rename to pkg/controller/jobs/trainjob/trainjob_multikueue_adapter.go
diff --git a/pkg/controller/jobs/kubeflow/trainjob/trainjob_webhook.go b/pkg/controller/jobs/trainjob/trainjob_webhook.go
similarity index 100%
rename from pkg/controller/jobs/kubeflow/trainjob/trainjob_webhook.go
rename to pkg/controller/jobs/trainjob/trainjob_webhook.go
diff --git a/site/go.mod b/site/go.mod
index fa341d783..cff84385d 100644
--- a/site/go.mod
+++ b/site/go.mod
@@ -1,5 +1,5 @@
 module sigs.k8s.io/kueue/site
 
-go 1.24
+go 1.25
 
 require github.com/google/docsy v0.11.0 // indirect



Test patch:
diff --git a/pkg/controller/jobs/kubeflow/trainjob/trainjob_controller_test.go b/pkg/controller/jobs/trainjob/trainjob_controller_test.go
similarity index 100%
rename from pkg/controller/jobs/kubeflow/trainjob/trainjob_controller_test.go
rename to pkg/controller/jobs/trainjob/trainjob_controller_test.go
diff --git a/pkg/controller/jobs/kubeflow/trainjob/trainjob_multikueue_adapter_test.go b/pkg/controller/jobs/trainjob/trainjob_multikueue_adapter_test.go
similarity index 100%
rename from pkg/controller/jobs/kubeflow/trainjob/trainjob_multikueue_adapter_test.go
rename to pkg/controller/jobs/trainjob/trainjob_multikueue_adapter_test.go
diff --git a/pkg/controller/jobs/kubeflow/trainjob/trainjob_webhook_test.go b/pkg/controller/jobs/trainjob/trainjob_webhook_test.go
similarity index 100%
rename from pkg/controller/jobs/kubeflow/trainjob/trainjob_webhook_test.go
rename to pkg/controller/jobs/trainjob/trainjob_webhook_test.go
diff --git a/test/e2e/multikueue/e2e_test.go b/test/e2e/multikueue/e2e_test.go
index 2e0bdbadd..83afe6d34 100644
--- a/test/e2e/multikueue/e2e_test.go
+++ b/test/e2e/multikueue/e2e_test.go
@@ -46,12 +46,12 @@ import (
 	workloadjob "sigs.k8s.io/kueue/pkg/controller/jobs/job"
 	workloadjobset "sigs.k8s.io/kueue/pkg/controller/jobs/jobset"
 	workloadpytorchjob "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/jobs/pytorchjob"
-	workloadtrainjob "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/trainjob"
 	workloadmpijob "sigs.k8s.io/kueue/pkg/controller/jobs/mpijob"
 	workloadpod "sigs.k8s.io/kueue/pkg/controller/jobs/pod"
 	podconstants "sigs.k8s.io/kueue/pkg/controller/jobs/pod/constants"
 	workloadraycluster "sigs.k8s.io/kueue/pkg/controller/jobs/raycluster"
 	workloadrayjob "sigs.k8s.io/kueue/pkg/controller/jobs/rayjob"
+	workloadtrainjob "sigs.k8s.io/kueue/pkg/controller/jobs/trainjob"
 	"sigs.k8s.io/kueue/pkg/util/admissioncheck"
 	utilpod "sigs.k8s.io/kueue/pkg/util/pod"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
diff --git a/test/e2e/singlecluster/trainjob_test.go b/test/e2e/singlecluster/trainjob_test.go
index db55b2279..67c91040b 100644
--- a/test/e2e/singlecluster/trainjob_test.go
+++ b/test/e2e/singlecluster/trainjob_test.go
@@ -26,7 +26,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
-	workloadtrainjob "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/trainjob"
+	workloadtrainjob "sigs.k8s.io/kueue/pkg/controller/jobs/trainjob"
 	"sigs.k8s.io/kueue/pkg/util/testing"
 	testingtrainjob "sigs.k8s.io/kueue/pkg/util/testingjobs/trainjob"
 	"sigs.k8s.io/kueue/test/util"
diff --git a/test/integration/multikueue/jobs_test.go b/test/integration/multikueue/jobs_test.go
index 4d314f217..7e2664a42 100644
--- a/test/integration/multikueue/jobs_test.go
+++ b/test/integration/multikueue/jobs_test.go
@@ -51,11 +51,11 @@ import (
 	workloadpytorchjob "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/jobs/pytorchjob"
 	workloadtfjob "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/jobs/tfjob"
 	workloadxgboostjob "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/jobs/xgboostjob"
-	workloadtrainjob "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/trainjob"
 	workloadmpijob "sigs.k8s.io/kueue/pkg/controller/jobs/mpijob"
 	workloadpod "sigs.k8s.io/kueue/pkg/controller/jobs/pod"
 	workloadraycluster "sigs.k8s.io/kueue/pkg/controller/jobs/raycluster"
 	workloadrayjob "sigs.k8s.io/kueue/pkg/controller/jobs/rayjob"
+	workloadtrainjob "sigs.k8s.io/kueue/pkg/controller/jobs/trainjob"
 	"sigs.k8s.io/kueue/pkg/features"
 	"sigs.k8s.io/kueue/pkg/util/admissioncheck"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
diff --git a/test/integration/multikueue/suite_test.go b/test/integration/multikueue/suite_test.go
index 14d0105dc..4ba15ee54 100644
--- a/test/integration/multikueue/suite_test.go
+++ b/test/integration/multikueue/suite_test.go
@@ -49,11 +49,11 @@ import (
 	workloadpytorchjob "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/jobs/pytorchjob"
 	workloadtfjob "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/jobs/tfjob"
 	workloadxgboostjob "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/jobs/xgboostjob"
-	workloadtrainjob "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/trainjob"
 	workloadmpijob "sigs.k8s.io/kueue/pkg/controller/jobs/mpijob"
 	workloadpod "sigs.k8s.io/kueue/pkg/controller/jobs/pod"
 	workloadraycluster "sigs.k8s.io/kueue/pkg/controller/jobs/raycluster"
 	workloadrayjob "sigs.k8s.io/kueue/pkg/controller/jobs/rayjob"
+	workloadtrainjob "sigs.k8s.io/kueue/pkg/controller/jobs/trainjob"
 	dispatcher "sigs.k8s.io/kueue/pkg/controller/workloaddispatcher"
 	"sigs.k8s.io/kueue/pkg/util/kubeversion"
 	utiltesting "sigs.k8s.io/kueue/pkg/util/testing"
diff --git a/test/integration/singlecluster/controller/jobs/trainjob/suite_test.go b/test/integration/singlecluster/controller/jobs/trainjob/suite_test.go
index 1dfd50e64..70018b80c 100644
--- a/test/integration/singlecluster/controller/jobs/trainjob/suite_test.go
+++ b/test/integration/singlecluster/controller/jobs/trainjob/suite_test.go
@@ -33,7 +33,7 @@ import (
 	"sigs.k8s.io/kueue/pkg/controller/core"
 	"sigs.k8s.io/kueue/pkg/controller/core/indexer"
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
-	"sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/trainjob"
+	"sigs.k8s.io/kueue/pkg/controller/jobs/trainjob"
 	"sigs.k8s.io/kueue/pkg/scheduler"
 	"sigs.k8s.io/kueue/test/integration/framework"
 	"sigs.k8s.io/kueue/test/util"
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 58c7bcfaf..df5da554b 100644
--- a/test/integration/singlecluster/controller/jobs/trainjob/trainjob_controller_test.go
+++ b/test/integration/singlecluster/controller/jobs/trainjob/trainjob_controller_test.go
@@ -33,7 +33,7 @@ import (
 	kueue "sigs.k8s.io/kueue/apis/kueue/v1beta1"
 	"sigs.k8s.io/kueue/pkg/controller/constants"
 	"sigs.k8s.io/kueue/pkg/controller/jobframework"
-	workloadtrainjob "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/trainjob"
+	workloadtrainjob "sigs.k8s.io/kueue/pkg/controller/jobs/trainjob"
 	"sigs.k8s.io/kueue/pkg/util/testing"
 	testingjobset "sigs.k8s.io/kueue/pkg/util/testingjobs/jobset"
 	testingtrainjob "sigs.k8s.io/kueue/pkg/util/testingjobs/trainjob"
diff --git a/test/integration/singlecluster/webhook/jobs/trainjob_webhook_test.go b/test/integration/singlecluster/webhook/jobs/trainjob_webhook_test.go
index 8e794170f..79d05bbf1 100644
--- a/test/integration/singlecluster/webhook/jobs/trainjob_webhook_test.go
+++ b/test/integration/singlecluster/webhook/jobs/trainjob_webhook_test.go
@@ -25,7 +25,7 @@ import (
 	"k8s.io/utils/ptr"
 
 	controllerconstants "sigs.k8s.io/kueue/pkg/controller/constants"
-	workloadtrainjob "sigs.k8s.io/kueue/pkg/controller/jobs/kubeflow/trainjob"
+	workloadtrainjob "sigs.k8s.io/kueue/pkg/controller/jobs/trainjob"
 	testingtrainjob "sigs.k8s.io/kueue/pkg/util/testingjobs/trainjob"
 	"sigs.k8s.io/kueue/test/util"
 )