Migrate InfoOptions.podSpecReplias and info.Scheduler.TotalRequests to info.TemplateSpec.PodSet #2524

tenzen-y · 2025-03-15T07:50:17Z

What this PR does / why we need it:
I removed duplicated internal data structure, InfoOptions.podSpecReplicas and Info.Scheduler.TotalRequests
and migrate its usage to Info.TemplateSpec.PodSet.

Which issue(s) this PR fixes (optional, in Fixes #<issue number>, #<issue number>, ... format, will close the issue(s) when PR gets merged):
Part-of #2495

Checklist:

Docs included if any changes are user facing

tenzen-y · 2025-03-15T07:58:36Z

pkg/runtime/framework/core/framework_test.go

-						},
-						constants.JobTrainerNode: {
-							Replicas:    100,         // Replicas is taken from TrainJob NumNodes.
-							PodRequests: resRequests, // TODO (andreyvelich): Add support for TrainJob ResourcesPerNode in TotalRequests.


Opened: #2525

tenzen-y · 2025-03-15T08:20:15Z

/assign @kubeflow/wg-training-leads @astefanutti

andreyvelich

Thank you for updating this @tenzen-y!
Mostly lgtm

andreyvelich · 2025-03-15T19:29:29Z

pkg/runtime/runtime.go

 	PodSets []PodSet
 }

 type PodSet struct {
+	// PodSet name is the name to identify PodSpec.
+	// This typically has the name stored in each PodSpec.
 	Name string
 	// If Name is trainer-node, CountForNonTrainer is null.
 	// For Trainer, PodSet Count should be stored in Info.RuntimePolicy.MLPolicy.NumNodes.
 	CountForNonTrainer *int32


Since we already removed the Trainer struct, do you want to refactor this count parameter in the next PR ?
I just feel like instead of updating the info.RuntimePolicy.MLPolicy.NumNodes, we can directly update count value in PodSet for the appropriate name (e.g. Launcher, Node):

trainer/pkg/runtime/framework/plugins/plainml/plainml.go

Line 56 in e627f78

info.RuntimePolicy.MLPolicy.NumNodes = numNodes

.

Which means PodSet.Count always dictates values for Parallelism and Completions.

I feel like it should simplifies a lot of logic we have.
For example:

trainer/pkg/runtime/framework/plugins/coscheduling/coscheduling.go

Lines 128 to 131 in e627f78

case constants.JobTrainerNode:

count = *info.RuntimePolicy.MLPolicy.NumNodes

default:

count = *ps.CountForNonTrainer

Yes, that's right. I will change MLPolocy data structure only for internal in the next PR.
I decoupled the changes from this PR since the internal MLPololicy data structure change has a lot of affection.

Yes, that's right. I will change MLPolocy data structure only for internal in the next PR.

Do you mean after we refactor MLPolicy for the Info object, we can change the Count value in PodSet ?

I will drop MLPolicy.NumNodes only from internal data structure, and rename PodSet.CoutNotForTrainer with PodSet.Count.

Yeah, that make sense!

andreyvelich · 2025-03-15T19:34:24Z

pkg/runtime/runtime.go

 	Name string
 	// If Name is trainer-node, CountForNonTrainer is null.
 	// For Trainer, PodSet Count should be stored in Info.RuntimePolicy.MLPolicy.NumNodes.
 	CountForNonTrainer *int32
+	InitContainers     []Container


Should we add InitContainers once we have a use-case when InitContainer can be configured via TrainJob ?

We want to keep InitContainers here since this is abstraction for Runtime.
So, each plugin can easily take and modify initContainers throughout cycle.

This is mostly preparation for overriding order changes (Plugin Change -> TrainJob -> Runtime).

Sounds good.

andreyvelich · 2025-03-15T19:45:40Z

pkg/runtime/runtime.go

 	Containers         []Container
 	Volumes            []corev1ac.VolumeApplyConfiguration
 	Endpoints          iter.Seq[string]
+	// The total PodSet requests can be calculated with
+	// SinglePodRequests x [CountForNonTrainer|RuntimePolicy.MLPolicy.NumNodes].
+	SinglePodRequests corev1.ResourceList


Is there any reason to store singlePodRequest rather than entire PodRequest ?
Since the only place we use it is in the co-scheduling plugin where we multiple it by count?

trainer/pkg/runtime/framework/plugins/coscheduling/coscheduling.go

Line 135 in e627f78

quantity.Mul(int64(count))

Because we want to avoid re-calculating every time we change numNodes.
If we have this SinglePodRequests, we can avoid calculating when we change numNodes in each plugin.

Sounds good, so you imagine scenarios when Count can be modified by multiple plugins ?

Yes after I drop MLPolicy.NumNodes.

…o info.TemplateSpec.PodSet Signed-off-by: Yuki Iwai <yuki.iwai.tz@gmail.com> # Conflicts: # pkg/runtime/core/trainingruntime.go # pkg/runtime/runtime.go

tenzen-y · 2025-03-16T10:07:09Z

I resolved conflicts.

andreyvelich

Thank you @tenzen-y!
/lgtm
/approve
feel free to unhold
/hold

google-oss-prow · 2025-03-16T16:56:18Z

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: andreyvelich

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~~ [andreyvelich]

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

tenzen-y · 2025-03-16T17:09:02Z

Thank you @tenzen-y! /lgtm /approve feel free to unhold /hold

Thank you
/hold cancel

google-oss-prow bot added the size/XL label Mar 15, 2025

google-oss-prow bot requested review from Electronic-Waste and kuizhiqing March 15, 2025 07:50

tenzen-y commented Mar 15, 2025

View reviewed changes

tenzen-y force-pushed the remove-pod-spec-replicas branch from c45ef80 to e627f78 Compare March 15, 2025 07:59

andreyvelich reviewed Mar 15, 2025

View reviewed changes

Migrate InfoOptions.podSpecReplias and info.Scheduler.TotalRequests t…

d0d0c08

…o info.TemplateSpec.PodSet Signed-off-by: Yuki Iwai <yuki.iwai.tz@gmail.com> # Conflicts: # pkg/runtime/core/trainingruntime.go # pkg/runtime/runtime.go

tenzen-y force-pushed the remove-pod-spec-replicas branch from e627f78 to d0d0c08 Compare March 16, 2025 10:05

andreyvelich reviewed Mar 16, 2025

View reviewed changes

google-oss-prow bot added the do-not-merge/hold label Mar 16, 2025

google-oss-prow bot assigned andreyvelich Mar 16, 2025

google-oss-prow bot added the lgtm label Mar 16, 2025

google-oss-prow bot added the approved label Mar 16, 2025

google-oss-prow bot removed the do-not-merge/hold label Mar 16, 2025

google-oss-prow bot merged commit 532e737 into kubeflow:master Mar 16, 2025
15 checks passed

tenzen-y deleted the remove-pod-spec-replicas branch March 16, 2025 17:10

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Migrate InfoOptions.podSpecReplias and info.Scheduler.TotalRequests to info.TemplateSpec.PodSet #2524

Migrate InfoOptions.podSpecReplias and info.Scheduler.TotalRequests to info.TemplateSpec.PodSet #2524

tenzen-y commented Mar 15, 2025 •

edited

Loading

tenzen-y Mar 15, 2025

tenzen-y commented Mar 15, 2025

andreyvelich left a comment

andreyvelich Mar 15, 2025

andreyvelich Mar 15, 2025

tenzen-y Mar 16, 2025

andreyvelich Mar 16, 2025

tenzen-y Mar 16, 2025

andreyvelich Mar 16, 2025

andreyvelich Mar 15, 2025

tenzen-y Mar 16, 2025

andreyvelich Mar 16, 2025

andreyvelich Mar 15, 2025

tenzen-y Mar 16, 2025

andreyvelich Mar 16, 2025

tenzen-y Mar 16, 2025

tenzen-y commented Mar 16, 2025

andreyvelich left a comment

google-oss-prow bot commented Mar 16, 2025

tenzen-y commented Mar 16, 2025

	case constants.JobTrainerNode:
	count = *info.RuntimePolicy.MLPolicy.NumNodes
	default:
	count = *ps.CountForNonTrainer

Migrate InfoOptions.podSpecReplias and info.Scheduler.TotalRequests to info.TemplateSpec.PodSet #2524

Migrate InfoOptions.podSpecReplias and info.Scheduler.TotalRequests to info.TemplateSpec.PodSet #2524

Conversation

tenzen-y commented Mar 15, 2025 • edited Loading

Choose a reason for hiding this comment

tenzen-y commented Mar 15, 2025

andreyvelich left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tenzen-y commented Mar 16, 2025

andreyvelich left a comment

Choose a reason for hiding this comment

google-oss-prow bot commented Mar 16, 2025

tenzen-y commented Mar 16, 2025

tenzen-y commented Mar 15, 2025 •

edited

Loading