Support JAX Runtimes

[Electronic-Waste]

What you would like to be added?

Part of: #2170

As we planned in the Kubeflow Trainer V2 API, we should support JAX runtime after we implement pytorch runtime.

The works include:

• Implement runtime plugins for JAX if necessary
• Create ClusterTrainingRuntime for JAX (single-node, multi-nodes)
• Add some unit tests and e2e tests
• Add some user guides for JAX runtimes

/area runtime
/cc @kubeflow/wg-training-leads @saileshd1402 @astefanutti @juliusvonkohout @franciscojavierarceo @varodrig @rareddy @thesuperzapper @seanlaii @deepanker13 @helenxie-bit @Doris-xm @truc0 @mahdikhashan

Why is this needed?

This is planned in KEP-2170.

Love this feature?

Give it a 👍 We prioritize the features with most 👍

[Electronic-Waste]

/remove-label lifecycle/needs-triage

[andreyvelich]

Thank you for creating this @Electronic-Waste!
I think, you can remove the KEP-2170 from the title since we can implement it after the first v2.0.0 release.
Also, should we split this task into two separate issues ?

We might require two KEPs for every Runtime, since it requires:

1. API changes
2. Extension of our runtime framework to support more plugins (or re-using PlainML plugin).
   WDYT @kubeflow/wg-training-leads @astefanutti @Electronic-Waste ?

[Electronic-Waste]

@andreyvelich SGTM. I'll split it into two seperate issues.

[tenzen-y]

• API changes
• Extension of our runtime framework to support more plugins (or re-using PlainML plugin).

Does "API" mean CRD? SDK API?

[andreyvelich]

Does "API" mean CRD? SDK API?

If we need to create a new MLPolicy: JAX, we need to extend the CRD APIs.
But I think we should also ask contributors to include SDK APIs changes (if that is needed) into KEP since this is the main interface for the end-users.

[tenzen-y]

Does "API" mean CRD? SDK API?

If we need to create a new MLPolicy: JAX, we need to extend the CRD APIs. But I think we should also ask contributors to include SDK APIs changes (if that is needed) into KEP since this is the main interface for the end-users.

SGTM

[Electronic-Waste]

@andreyvelich @tenzen-y Shall we create two GSoC projects for supporting JAX/TF Runtimes? They need two separate KEPs.

[andreyvelich]

@andreyvelich @tenzen-y Shall we create two GSoC projects for supporting JAX/TF Runtimes? They need two separate KEPs.

I am not sure if we have sufficient number of slots since we already propose 12 projects.
If you have bandwidth to mentor more students, maybe we could separate it.

[Electronic-Waste]

@andreyvelich This is the first time for me to mentor students in GSoC. And I will serve as primary mentor for 2 projects and backup mentor for some others now. I'm not sure whether I could handle 3 projects... Anyway, let's discuss it in the upcoming WG Training/AutoML Call.