KEP-2170: Implement JobSet, PlainML, and Torch Plugins

[andreyvelich]

Part of: #2290.

I implemented JobSet and PlainML plugins, and I was able to submit our first TrainJob and ClusterTrainingRuntime 🎉

Working ClusterTrainingRuntime

apiVersion: kubeflow.org/v2alpha1
kind: ClusterTrainingRuntime
metadata:
  name: first-runtime
spec:
  mlPolicy:
    numNodes: 5
  template:
    spec:
      replicatedJobs:
        - name: trainer-node
          template:
            spec:
              template:
                spec:
                  containers:
                    - name: trainer
                      image: python:alpine3.20
                      command:
                        - "/bin/sh"
                        - -c
                        - "echo 'Hello From Training Runtime'"
---
apiVersion: kubeflow.org/v2alpha1
kind: TrainJob
metadata:
  name: first-job
spec:
  runtimeRef:
    name: first-runtime
    kind: ClusterTrainingRuntime
    apiGroup: kubeflow.org
  trainer:
    numNodes: 2

A few details:

• We discussed with @tenzen-y that Info object doesn't need to have Obj class, since we can directly pass the JobTemplateSpec into RunComponentBuilderPlugins() API.
• I added the Trainer struct in Info object which represents values combined from MLPolicy and TrainJob. Those values will be inserted into JobSet.
• I need to update the unit tests in this PR.
• I need to implement torch plugin.
• I need to implement env variables enforcement.

/assign @kubeflow/wg-training-leads @varshaprasad96 @akshaychitneni @deepanker13 @helenxie-bit @Electronic-Waste @saileshd1402 @kannon92 @kuizhiqing @shravan-achar

[andreyvelich]

/hold for review

[coveralls]

Pull Request Test Coverage Report for Build 11617265788

Details

• 0 of 0 changed or added relevant lines in 0 files are covered.
• No unchanged relevant lines lost coverage.
• Overall coverage remained the same at 100.0%

---

Totals
Change from base Build 11542660312:	0.0%
Covered Lines:	77
Relevant Lines:	77

---

💛 - Coveralls

[andreyvelich]

Torch Runtime is working 🎉

Simple Torch runtime

apiVersion: kubeflow.org/v2alpha1
kind: ClusterTrainingRuntime
metadata:
  name: torch-runtime
spec:
  mlPolicy:
    numNodes: 2
    torch:
      numProcPerNode: "1"
  template:
    spec:
      replicatedJobs:
        - name: trainer-node
          template:
            spec:
              template:
                spec:
                  containers:
                    - name: trainer
                      image: docker.io/andreyvelichkevich/test-torch
---
apiVersion: kubeflow.org/v2alpha1
kind: TrainJob
metadata:
  name: first-job
spec:
  runtimeRef:
    name: torch-runtime
    kind: ClusterTrainingRuntime
    apiGroup: kubeflow.org
  trainer:
    numNodes: 3
    numProcPerNode: "5"
    env:
      - name: TRAIN_JOB
        value: env

[tenzen-y]

Sorry for the delay review.
I did not review any testings yet.

[tenzen-y]

Do we still need the Policy?

[andreyvelich]

Actually, I am still getting some data from Policy in the EnforceMLPolicy:

training-operator/pkg/runtime.v2/framework/plugins/plainml/plainml.go

Line 50 in f00c1d3

numNodes := info.MLPolicy.NumNodes

Do you think, that we should just remove Policy API from Info object and bypass the TrainingRuntime to the enforcement APIs ?

[andreyvelich]

I refactored the Info object to remove Policies:

training-operator/pkg/runtime.v2/runtime.go

Lines 26 to 31 in e5165cb

	type Info struct {
	Labels map[string]string
	Annotations map[string]string
	Trainer
	*Scheduler
	}

.
Right now, we pass MLPolicy and PodGroupPolicy spec from the Runtime directly to the Enforce() APIs.

[tenzen-y]

I think that the Trainer needs to have scheduler since we need to create the PodGroup and do some else scheduling logic.
If we create the PodGroup against a whole of TrainJob, that will not schedule to any Node for ever. Because we need to start the initializer before the Trainer.

[andreyvelich]

That makes sense.
Does it mean that we should apply the coscheduling PodLabel only to the Trainer Job, not to the Initializer Job ?

[tenzen-y]

Does it mean that we should apply the coscheduling PodLabel only to the Trainer Job, not to the Initializer Job ?

Yes. When users want to take the intializers separate from trainer, they want to specify another PodGroup name to initializers.

[andreyvelich]

I see, do you want to update this change in the followup PR once I add the builder for Dataset and Model ?

[tenzen-y]

I see, do you want to update this change in the followup PR once I add the builder for Dataset and Model ?

This can be followed up.
I will take it during you impl Datase initializer since this is obviously regressions as I mentioned the above.
The current impl bring us a whole of Job will be stuck forever.

[akshaychitneni]

I think we should have a utility function to calculate the quantity passing in the obj

[andreyvelich]

We already have it here to calculate PodRequests based on the PodSpec: https://github.com/kubeflow/training-operator/blob/master/pkg/runtime.v2/runtime.go#L121-L124.

Do you mean having another utility just for the unit tests ?

[akshaychitneni]

yes utility to compute value in tests

[tenzen-y]

I was not sure the reason why @akshaychitneni requests to implement the util.

[andreyvelich]

@akshaychitneni Please can you give an example of code snippet that you want to use instead of:

corev1.ResourceCPU: resource.MustParse("31")

[akshaychitneni]

I was referring to calculating 31. Probably it is simpler as it is (numNodes + 1) * (Num of CPU configured in the runtime) in this case

[andreyvelich]

I see, I think it depends on the unit test. For example, we have unit tests where we add CPU + Memory for the Pod requests:

training-operator/pkg/runtime.v2/framework/core/framework_test.go

Lines 359 to 360 in 85d0cf1

	corev1.ResourceCPU: resource.MustParse("1"),
	corev1.ResourceMemory: resource.MustParse("4Gi"),

Additionally, in the future we might want to define different resources for initializer containers and trainer to verify that we correctly calculate MinResource for the PodGroup.

Let me add the TODO statement to see if we can create helper function for our tests to calculate these values based on number of replicas and resources per replica.

[tenzen-y]

Actually, container resource computation is not easy, we need to consider runtimeClass (not training runtime, it's core API runtimeClass), limitRange, and initContainer and so on. So, I would like to keep manually resource creation for the testing.

If we provide such utils with complex computing, the lib could potentially bug, and we will miss the bugs in main testing cases.

[tenzen-y]

I submitted the production code review comments, first.

[tenzen-y]

We do not want to surface the mlPlicy to the JobPipeline framework directly since the framework is not only for the TrainingRuntime.
So, could you move the mlPolicy propagation into the runtime initialization as we disscussed before?

[tenzen-y]

Suggested change

	if info == nil || info.Scheduler == nil || info.Scheduler.ScheduleTimeoutSeconds == nil || trainJob == nil {
	if info == nil || info.Scheduler == nil || trainJob == nil {

As we discussed berore, they need to specify the empty cosheduler parameters when they do not want to modify the schedulerTimeoutSeconds, right?

[andreyvelich]

Do we really want to allow this ? Otherwise, runtime might look as follows:

spec:
  mlPolicy:
    numNodes: 10
    torch:
  podGroupPolicy:
    coscheduling:

Is it a valid object from the API perspective ?

[tenzen-y]

As we discussed before, they can specify the empty object like this:

spec:
  mlPolicy:
    numNodes: 10
    torch:
  podGroupPolicy:
    coscheduling: {}

[andreyvelich]

I reverted it back to your original condition check:

training-operator/pkg/runtime.v2/framework/plugins/coscheduling/coscheduling.go

Line 92 in 894cded

if info == nil || info.RuntimePolicy.PodGroupPolicy == nil || trainJob == nil {

[tenzen-y]

I meant that the reason why we do not store thre totalRequests for the initializers?

[tenzen-y]

Comments for the unit testings.

[tenzen-y]

I think that the Trainer needs to have scheduler since we need to create the PodGroup and do some else scheduling logic.
If we create the PodGroup against a whole of TrainJob, that will not schedule to any Node for ever. Because we need to start the initializer before the Trainer.

[tenzen-y]

I was not sure the reason why @akshaychitneni requests to implement the util.

[tenzen-y]

Comments for the Integration testings

[tenzen-y]

Could you add a case if we can build and deploy Torch JobSet like ginkgo.It("Should succeed to create TrainJob with Torch TrainingRuntime", func() {

[deepanker13]

ctx context is not getting used anywhere throughout the code, is it necessary to pass?

[andreyvelich]

I think, we use it here to get the existing object and compare it
https://github.com/kubeflow/training-operator/blob/master/pkg/runtime.v2/framework/plugins/jobset/jobset.go#L82-L84

[andreyvelich]

I keep the original runtime policies values in the Info object for now, since we can't pass the specific policies to the EnforceMLPolicy and EnforcePodGroupPolicy() APIs. Similar to what we do today:

training-operator/pkg/runtime.v2/runtime.go

Line 40 in 62a058d

Policy

In the future, we will discuss how we can optimize the Info object structure.

cc @tenzen-y

[tenzen-y]

Yeah, that is at least a requirement. Let's revisit here later. I can take that in the next week.

[tenzen-y]

What is the reason that we forcefully override the Pod labels?

[andreyvelich]

We only use PodLabels in the Scheduler right now, that is why we just creating a new Map here.
In the future, if we should do any Pod labels merge for the Info object, we should revisit it.

[tenzen-y]

Basically, this is pluggabble mechanism, so we should consider the other plugin.
So, as before initialization impl, we should initialize map only when the pod label is emtpy

[andreyvelich]

Alright, let me change it.

[tenzen-y]

Thank you for this effort!
Let's open some follow-up issue!
/lgtm
/apprve

[tenzen-y]

/approve

[tenzen-y]

/hold cancel

[google-oss-prow]

[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