KEP-2170: Support hundreds and thousands worker nodes for a single training Job

[tenzen-y]

What you would like to be added?

We should support the multiple replicas per a replicatedJob like:

[...]
spec:
  replicatedJobs:
  - name:
    replicas: 5
[...]

Why is this needed?

Currently, we enforce 1 to the JobSet ReplicatedJob replicas:

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

Lines 108 to 110 in 9e46f9d

	for _, rJob := range jobSetTemplateSpec.Spec.ReplicatedJobs {
	// By default every ReplicatedJob has only 1 replica.
	opts = append(opts, runtime.WithPodSpecReplicas(rJob.Name, 1, rJob.Template.Spec.Template.Spec))

However, when the single worker replicatedJob has batch/v1 Job with hundreds and thousands of completions (.spec.completions), this brings us a significant reconciling delay since the job-controller (combined within kube-controller-manager) reconciliation will take much longer time due to thousands of Pods, then following Jobs will be stuck in the workqueue.

spec:
  replicatedJobs:
  - name: training-node
    replicas: 1
    template:
      spec:
        completions: 2000
        parallelism: 2000

After that, the kube-controller-manger workqueue depth will be much deeper, which could potentially cause a memory leak.
Finally, the kube-controller-manager continues to restart, and any kind of Workload (even StatefulSet and Deployment) will fall unhandled.

Love this feature?

Give it a 👍 We prioritize the features with most 👍

[tenzen-y]

/remove-label lifecycle/needs-triage

[AydanPirani]

/assign

I can work on this!

@tenzen-y Quick question - I was looking at runtime.go#L71, and it seems that we already have support for PodSpecReplicas.

Does that mean that the majority of the implementation will be on the parsing side of things? (Adding a parameter for job name)

Of course, this is all assuming that there is already support to run multiple replicas - can you confirm this?

Thanks in advance!