feat(api): Sync TrainJob JobsStatus from JobSet ReplicatedJobsStatus #2802

astefanutti · 2025-08-28T11:37:08Z

What this PR does / why we need it:

This PR adds the JobsStatusPlugin runtime interface and implements it in the JobSet plugin so the TrainJob JobsStatus is synced from JobSet ReplicatedJobsStatus.

/cc @kubeflow/kubeflow-trainer-team @toVersus

Checklist:

Docs included if any changes are user facing

coveralls · 2025-08-28T11:41:39Z

Pull Request Test Coverage Report for Build 18313170282

Details

8 of 38 (21.05%) changed or added relevant lines in 5 files are covered.
3 unchanged lines in 3 files lost coverage.
Overall coverage decreased (-0.3%) to 54.234%

Changes Missing Coverage	Changed/Added Lines	%
pkg/runtime/core/clustertrainingruntime.go	2	0.0%
pkg/runtime/core/trainingruntime.go	2	0.0%
pkg/controller/trainjob_controller.go	7	0.0%
pkg/runtime/framework/plugins/jobset/jobset.go	19	0.0%

Files with Coverage Reduction	New Missed Lines	%
pkg/controller/trainjob_controller.go	1	0.0%
pkg/runtime/core/clustertrainingruntime.go	1	44.44%
pkg/runtime/core/trainingruntime.go	1	69.83%

Totals
Change from base Build 18271389532:	-0.3%
Covered Lines:	1204
Relevant Lines:	2220

💛 - Coveralls

andreyvelich

Thank you for this @astefanutti!
Please can you also add the integration: https://github.com/kubeflow/trainer/blob/master/test/integration/controller/trainjob_controller_test.go#L74 and e2e test: https://github.com/kubeflow/trainer/blob/master/test/e2e/e2e_test.go#L67, where we check that JobsStatus.

pkg/controller/trainjob_controller.go

andreyvelich · 2025-08-29T01:31:57Z

pkg/runtime/framework/core/framework.go

+	}
+	if len(f.jobsStatusPlugins) != 0 {
+		return f.jobsStatusPlugins[0].JobsStatus(ctx, trainJob)


@tenzen-y @astefanutti What is the use-case when we can have more than one terminalConditionPlugin or jobsStatusPlugin ?

@andreyvelich that would be if two plugins implement the same plugin interface so they would conflict.

The check is currently performed for each call, maybe we could do that during the runtime framework initialization.

But I am trying to understand when users can define more than 1 implementation for such plugins ?
Terminal condition or job status should always be defined by one plugin, unless I am missing something ?

Conceptually yes, but practically it is possible to have multiple plugins implementing the same interface like framework.TerminalConditionPlugin.

We could imagine this happens if we had other backend than JobSet, e.g., LWS and both would be made active by "mistake".

both would be made active by "mistake".

If that is set by mistake, shall we validate that only one plugin can be activated for such interfaces during framework initialization ?

Also, make those plugins as not arrays here:

trainer/pkg/runtime/framework/core/framework.go

Lines 48 to 49 in 389023c

terminalConditionPlugins []framework.TerminalConditionPlugin

jobsStatusPlugins []framework.JobsStatusPlugin

?

If that is set by mistake, shall we validate that only one plugin can be activated for such interfaces during framework initialization ?

I agree, that's what I meant when I said "maybe we could do that during the runtime framework initialization.".

@andreyvelich I've updated the logic for those "singleton" plugins so it fails-fast on at init time.

astefanutti · 2025-09-01T16:15:57Z

Thank you for this @astefanutti! Please can you also add the integration: https://github.com/kubeflow/trainer/blob/master/test/integration/controller/trainjob_controller_test.go#L74 and e2e test: https://github.com/kubeflow/trainer/blob/master/test/e2e/e2e_test.go#L67, where we check that JobsStatus.

@andreyvelich I've added/updated integration and e2e tests. PTAL when you can.

pkg/controller/trainjob_controller.go

andreyvelich · 2025-09-01T16:54:27Z

pkg/runtime/framework/core/framework.go

+	}
+	if len(f.jobsStatusPlugins) != 0 {
+		return f.jobsStatusPlugins[0].JobsStatus(ctx, trainJob)


Also, make those plugins as not arrays here:

trainer/pkg/runtime/framework/core/framework.go

Lines 48 to 49 in 389023c

terminalConditionPlugins []framework.TerminalConditionPlugin

jobsStatusPlugins []framework.JobsStatusPlugin

?

andreyvelich · 2025-09-01T16:57:09Z

pkg/runtime/framework/plugins/jobset/jobset.go

+			Name:      status.Name,
+			Ready:     status.Ready,
+			Succeeded: status.Succeeded,
+			Failed:    status.Failed,
+			Active:    status.Active,
+			Suspended: status.Suspended,


I am wondering whether we will introduce a breaking change after this issue: kubernetes-sigs/jobset#723 ?
I proposed that we don't use Succeeded as Job condition in rJob.
Any thoughts @astefanutti @tenzen-y @kannon92 @kubeflow/kubeflow-trainer-team ?

I agree this would be better to be consistent across projects, thought I'd suggest not to hold this PR as it'll be a breaking change anyway and have it addressed separately so we can clearly "release-note" it as a breaking change. WDYT?

As we discussed offline in Slack, we will use Succeeded condition for now.
We will change this API to Complete in a future version, once JobSet migrates its APIs.

astefanutti · 2025-09-03T07:47:40Z

/retest

andreyvelich · 2025-09-26T13:53:37Z

/milestone v2.1

andreyvelich · 2025-09-26T15:54:39Z

pkg/runtime/framework/core/framework.go

+	errorTooManyTerminalConditionPlugin = errors.New("too many TerminalCondition plugins are registered")
+	errorTooManyJobsStatusPlugin        = errors.New("too many JobsStatus plugins are registered")


Since we can only register one plugin for TerminalCondition and JobsStatus, do we need these errors ?

That's actually the error that's returned at init time in case more than one plugin is being registered.

But it won't be possible since we initialize the Framework here:

trainer/pkg/runtime/framework/core/framework.go

Lines 53 to 55 in 848a4bd

f := &Framework{

registry: r,

}

So this condition will never be executed:

trainer/pkg/runtime/framework/core/framework.go

Lines 82 to 85 in 848a4bd

if p, ok := plugin.(framework.TerminalConditionPlugin); ok {

if f.terminalConditionPlugin != nil {

return nil, errorTooManyTerminalConditionPlugin

}

Am I missing something ?

Right and that's where the errors are returned.

@astefanutti Can you explain the flow when f.terminalConditionPlugin != nil when we execute framework's New() function ?

@tenzen-y I am wondering if there is any need to distinguish plugins that we define in the Registry and plugins that are part of those "registered" plugins (e.g. TerminalCondition plugin is part of JobSet plugin).

I could not catch the reason why you want to do that.
I think the current @astefanutti 's implementations are a straightforward way.

#2802 (comment)

@tenzen-y I was asking what does "Plugin" mean in the Extension Framework from your point of view ?
And how are we going to allow users to extend the framework with their own plugins ?

Ah, got it. The trainer provides the 2 level abstraction layers which is Runtime Framework and Plugin.
But, sometimes, the plugin is only for a specific Runtime Framework like the JobSet plugin.

I guess that we might want to have mechanisms to represent the tie between the Runtime Framework and the Plugin, but it would be better to avoid adding any validation relationship to allow them to reuse the implemented Plugin wherever the Runtime Framework.

Anyway, we can consider that as an another enhancement.

Sounds good.

Signed-off-by: Antonin Stefanutti <antonin@stefanutti.fr>

tenzen-y · 2025-10-06T15:48:11Z

pkg/runtime/framework/core/framework.go

+	errorTooManyTerminalConditionPlugin = errors.New("too many TerminalCondition plugins are registered")
+	errorTooManyJobsStatusPlugin        = errors.New("too many JobsStatus plugins are registered")


@tenzen-y I am wondering if there is any need to distinguish plugins that we define in the Registry and plugins that are part of those "registered" plugins (e.g. TerminalCondition plugin is part of JobSet plugin).

I could not catch the reason why you want to do that.
I think the current @astefanutti 's implementations are a straightforward way.

tenzen-y · 2025-10-06T15:56:16Z

pkg/runtime/interface.go

 	TerminalCondition(ctx context.Context, trainJob *trainer.TrainJob) (*metav1.Condition, error)
+	JobsStatus(ctx context.Context, trainJob *trainer.TrainJob) ([]trainer.JobStatus, error)


Suggested change

TerminalCondition(ctx context.Context, trainJob *trainer.TrainJob) (*metav1.Condition, error)

JobsStatus(ctx context.Context, trainJob *trainer.TrainJob) ([]trainer.JobStatus, error)

UnderlyingJobState(ctx context.Context, trainJob *trainer.TrainJob) (*metav1.Condition, []trainer.JobStatus, error)

I think that TerminalCondition and JobsStatus responsibilities are almost same which is just propagate the Job state from underlying one to TrainJob.

Make sense for me to consolidate those two plugins together.

I've folded the two plugins together. PTAL.

tenzen-y · 2025-10-06T15:56:39Z

pkg/runtime/framework/interface.go

 type TerminalConditionPlugin interface {
 	Plugin
 	TerminalCondition(ctx context.Context, trainJob *trainer.TrainJob) (*metav1.Condition, error)
 }
+
+type JobsStatusPlugin interface {
+	Plugin
+	JobsStatus(ctx context.Context, trainJob *trainer.TrainJob) ([]trainer.JobStatus, error)
+}


ditto related to consolidation.

Updated with the two plugins together. PTAL.

tenzen-y · 2025-10-06T15:57:20Z

pkg/runtime/framework/core/framework.go

-	// TODO (tenzen-y): Once we provide the Configuration API, we should validate which plugin should have terminalCondition execution points.
-	if len(f.terminalConditionPlugins) > 1 {


Any reason why we can remove the verification of the number of plugins?

As we discussed here, users won't be able to register more than one plugin for terminal condition: #2802 (comment)

Yes, the check is now moved at initialization, rather than at (first) execution.

Ah, right. The new way makes more sense.
Thank you!

Signed-off-by: Antonin Stefanutti <antonin@stefanutti.fr>

tenzen-y

LGTM
Thank you!
/lgtm
/approve

google-oss-prow · 2025-10-07T13:18:08Z

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: tenzen-y

The full list of commands accepted by this bot can be found here.

The pull request process is described here

Needs approval from an approver in each of these files:

~~OWNERS~~ [tenzen-y]

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

google-oss-prow bot requested review from jinchihe and kuizhiqing August 28, 2025 11:37

google-oss-prow bot added the size/L label Aug 28, 2025

astefanutti mentioned this pull request Aug 28, 2025

feat: support add pending, runing condition to trainjob #2725

Open

1 task

andreyvelich reviewed Aug 29, 2025

View reviewed changes

astefanutti force-pushed the pr-25 branch from 5863d08 to c20ef0d Compare September 1, 2025 08:29

google-oss-prow bot added size/XL and removed size/L labels Sep 1, 2025

astefanutti force-pushed the pr-25 branch from 87c9fa9 to f6a25f8 Compare September 1, 2025 10:19

andreyvelich reviewed Sep 1, 2025

View reviewed changes

astefanutti force-pushed the pr-25 branch 2 times, most recently from c25b925 to 3e4fc1a Compare September 2, 2025 07:18

astefanutti force-pushed the pr-25 branch from 3e4fc1a to 848a4bd Compare September 8, 2025 07:49

google-oss-prow bot added this to the v2.1 milestone Sep 26, 2025

andreyvelich mentioned this pull request Sep 26, 2025

Kubeflow Trainer v2.1 Release #2859

Open

7 tasks

andreyvelich reviewed Sep 26, 2025

View reviewed changes

astefanutti added 7 commits October 6, 2025 15:16

feat(api): Sync TrainJob JobsStatus from JobSet ReplicatedJobsStatus

223f1b5

Signed-off-by: Antonin Stefanutti <antonin@stefanutti.fr>

Add integration tests

783630d

Signed-off-by: Antonin Stefanutti <antonin@stefanutti.fr>

Update e2e tests

bfe4234

Signed-off-by: Antonin Stefanutti <antonin@stefanutti.fr>

Remove extra check

a1287d3

Signed-off-by: Antonin Stefanutti <antonin@stefanutti.fr>

Sort JobsStatus in e2e tests

efc5e9d

Signed-off-by: Antonin Stefanutti <antonin@stefanutti.fr>

Fix e2e test for MPI job

ab46fa6

Signed-off-by: Antonin Stefanutti <antonin@stefanutti.fr>

Fail-fast when multiple terminal condition and JobsStatus plugins exist

fc55beb

Signed-off-by: Antonin Stefanutti <antonin@stefanutti.fr>

astefanutti force-pushed the pr-25 branch from 848a4bd to fc55beb Compare October 6, 2025 13:16

tenzen-y reviewed Oct 6, 2025

View reviewed changes

Fold TerminalCondition and JobsStatus plugins

7d34b6e

Signed-off-by: Antonin Stefanutti <antonin@stefanutti.fr>

tenzen-y reviewed Oct 7, 2025

View reviewed changes

google-oss-prow bot assigned tenzen-y Oct 7, 2025

google-oss-prow bot added the lgtm label Oct 7, 2025

google-oss-prow bot added the approved label Oct 7, 2025

google-oss-prow bot merged commit c1bc94d into kubeflow:master Oct 7, 2025
26 of 27 checks passed

astefanutti deleted the pr-25 branch October 7, 2025 13:23

	terminalConditionPlugins []framework.TerminalConditionPlugin
	jobsStatusPlugins []framework.JobsStatusPlugin

		errorTooManyTerminalConditionPlugin = errors.New("too many TerminalCondition plugins are registered")
		errorTooManyJobsStatusPlugin = errors.New("too many JobsStatus plugins are registered")

	if p, ok := plugin.(framework.TerminalConditionPlugin); ok {
	if f.terminalConditionPlugin != nil {
	return nil, errorTooManyTerminalConditionPlugin
	}

		TerminalCondition(ctx context.Context, trainJob trainer.TrainJob) (metav1.Condition, error)
		JobsStatus(ctx context.Context, trainJob *trainer.TrainJob) ([]trainer.JobStatus, error)

	TerminalCondition(ctx context.Context, trainJob trainer.TrainJob) (metav1.Condition, error)
	JobsStatus(ctx context.Context, trainJob *trainer.TrainJob) ([]trainer.JobStatus, error)
	UnderlyingJobState(ctx context.Context, trainJob trainer.TrainJob) (metav1.Condition, []trainer.JobStatus, error)

		// TODO (tenzen-y): Once we provide the Configuration API, we should validate which plugin should have terminalCondition execution points.
		if len(f.terminalConditionPlugins) > 1 {

feat(api): Sync TrainJob JobsStatus from JobSet ReplicatedJobsStatus #2802

feat(api): Sync TrainJob JobsStatus from JobSet ReplicatedJobsStatus #2802

Uh oh!

Conversation

astefanutti commented Aug 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

coveralls commented Aug 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pull Request Test Coverage Report for Build 18313170282

Details

💛 - Coveralls

Uh oh!

andreyvelich left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

astefanutti commented Sep 1, 2025

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

astefanutti Sep 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

astefanutti commented Sep 3, 2025

Uh oh!

andreyvelich commented Sep 26, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

astefanutti commented Aug 28, 2025 •

edited

Loading

coveralls commented Aug 28, 2025 •

edited

Loading

astefanutti Sep 2, 2025 •

edited

Loading