Add JAX API #2163

sandipanpanda · 2024-07-10T21:02:26Z

What this PR does / why we need it:
Integrate JAX with Kubeflow Training Operator to support JAXJob

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

Checklist:

Docs included if any changes are user facing

sandipanpanda · 2024-07-10T21:03:30Z

/cc @andreyvelich @tenzen-y

sandipanpanda · 2024-07-10T21:05:01Z

/area gsoc

coveralls · 2024-07-10T21:21:37Z

Pull Request Test Coverage Report for Build 9912856908

Warning: This coverage report may be inaccurate.

This pull request's base commit is no longer the HEAD commit of its target branch. This means it includes changes from outside the original pull request, including, potentially, unrelated coverage changes.

For more information on this, see Tracking coverage changes with pull request builds.
To avoid this issue with future PRs, see these Recommended CI Configurations.
For a quick fix, rebase this PR at GitHub. Your next report should be accurate.

Details

4 of 495 (0.81%) changed or added relevant lines in 13 files are covered.
1 unchanged line in 1 file lost coverage.
Overall coverage decreased (-1.3%) to 34.082%

Changes Missing Coverage	Changed/Added Lines	%
pkg/client/informers/externalversions/generic.go	2	0.0%
pkg/client/clientset/versioned/typed/kubeflow.org/v1/fake/fake_kubeflow.org_client.go	3	0.0%
pkg/client/clientset/versioned/typed/kubeflow.org/v1/kubeflow.org_client.go	3	0.0%
pkg/client/informers/externalversions/kubeflow.org/v1/interface.go	3	0.0%
pkg/apis/kubeflow.org/v1/zz_generated.defaults.go	10	0.0%
pkg/client/informers/externalversions/kubeflow.org/v1/jaxjob.go	25	0.0%
pkg/client/listers/kubeflow.org/v1/jaxjob.go	25	0.0%
pkg/apis/kubeflow.org/v1/jax_defaults.go	30	0.0%
pkg/apis/kubeflow.org/v1/zz_generated.deepcopy.go	65	0.0%
pkg/client/clientset/versioned/typed/kubeflow.org/v1/fake/fake_jaxjob.go	74	0.0%

Files with Coverage Reduction	New Missed Lines	%
pkg/controller.v1/mpi/mpijob.go	1	91.06%

Totals
Change from base Build 9906458669:	-1.3%
Covered Lines:	4383
Relevant Lines:	12860

💛 - Coveralls

tenzen-y

Almost lgtm

tenzen-y · 2024-07-12T15:42:54Z

pkg/apis/kubeflow.org/v1/jax_defaults.go

+	"k8s.io/apimachinery/pkg/runtime"
+)
+
+func addJAXJobDefaultingFuncs(scheme *runtime.Scheme) error {


Suggested change

func addJAXJobDefaultingFuncs(scheme *runtime.Scheme) error {

func addJAXDefaultingFuncs(scheme *runtime.Scheme) error {

We usually do not add the Job suffix in most frameworks, although the XGBoost has the prefix.
Let's keep it simple.

tenzen-y · 2024-07-12T15:43:46Z

pkg/apis/kubeflow.org/v1/jax_defaults.go

+}
+
+// setJAXJobDefaultPort sets the default ports for jax container.
+func setJAXJobDefaultPort(spec *corev1.PodSpec) {


Could you eliminate the suffix, Job in all functions in this defaulter file?

pkg/apis/kubeflow.org/v1/jax_defaults.go

andreyvelich

Thank you for this @sandipanpanda!

andreyvelich · 2024-07-12T18:56:21Z

pkg/apis/kubeflow.org/v1/jax_types.go

+	// JAXJobFrameworkName is the name of the ML Framework
+	JAXJobFrameworkName = "jax"
+	// JAXJobReplicaTypeCoordinator is the type of Coordinator of distributed JAX
+	JAXJobReplicaTypeCoordinator ReplicaType = "Coordinator"


As we discussed today, we don't need to have separate replica for JAX, since Pod template will be the same and Worker with index 0 will the Coordinator.
Does it sound good to you @tenzen-y ?

I'm ok with v1 API since we can not apply separate runPolicies only to coordinator.
In the v2 API, I believe that dedicated replica for the coordinator would be worth it since we often want to treat the coordinator runPolicies.

But, v2 API is a separate discussion.

Makes sense @tenzen-y. Do you think, setting various runPolicy for JAX Coordinator and Workers makes sense even that JAX follows SPMD model: https://jax.readthedocs.io/en/latest/multi_process.html#launching-jax-processes ?
E.g. all JAX hosts should be always running.

In the v1 Training API, we don't have any ability to configure runPolicy on each replicas (coordinator and worker).
So, let's discuss if we should add the Coordinator role in the v2 API.

Support JAXJob Signed-off-by: Sandipan Panda <[email protected]>

sandipanpanda · 2024-07-12T19:07:20Z

I have updated the PR. PTAL.

terrytangyuan

LGTM to me. Pretty straightforward. Looking forward to the controller implementation!

tenzen-y · 2024-08-01T18:17:51Z

@sandipanpanda Is this still working in progress?

sandipanpanda · 2024-08-01T18:56:55Z

@sandipanpanda Is this still working in progress?

@tenzen-y I need to add the controller part, should that be in a separate PR?

andreyvelich · 2024-08-01T19:57:05Z

@sandipanpanda Is this still working in progress?

@tenzen-y I need to add the controller part, should that be in a separate PR?

Yes, I think, we should implement the controller logic as followup PRs.
@kubeflow/wg-training-leads @shravan-achar Do you have any other comments on the JAXJob API design ?

tenzen-y · 2024-08-02T03:57:32Z

Yes, I imagined that the controller and examples are implemented in the other PRs.

/lgtm
/approve

google-oss-prow · 2024-08-02T03:57:38Z

[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 added the do-not-merge/work-in-progress label Jul 10, 2024

google-oss-prow bot requested review from jinchihe and kuizhiqing July 10, 2024 21:02

google-oss-prow bot added the size/XXL label Jul 10, 2024

google-oss-prow bot requested review from andreyvelich and tenzen-y July 10, 2024 21:04

google-oss-prow bot added the area/gsoc label Jul 10, 2024

tenzen-y reviewed Jul 12, 2024

View reviewed changes

andreyvelich reviewed Jul 12, 2024

View reviewed changes

Add JAX API

3ac065f

Support JAXJob Signed-off-by: Sandipan Panda <[email protected]>

sandipanpanda force-pushed the jaxapi branch from 95edc8e to 3ac065f Compare July 12, 2024 19:04

sandipanpanda changed the title ~~[WIP] Add JAX API~~ Add JAX API Jul 24, 2024

google-oss-prow bot removed the do-not-merge/work-in-progress label Jul 24, 2024

sandipanpanda changed the title ~~Add JAX API~~ [WIP] Add JAX API Jul 24, 2024

google-oss-prow bot added the do-not-merge/work-in-progress label Jul 24, 2024

terrytangyuan reviewed Jul 26, 2024

View reviewed changes

sandipanpanda changed the title ~~[WIP] Add JAX API~~ Add JAX API Aug 1, 2024

google-oss-prow bot removed the do-not-merge/work-in-progress label Aug 1, 2024

google-oss-prow bot assigned tenzen-y Aug 2, 2024

google-oss-prow bot added the lgtm label Aug 2, 2024

google-oss-prow bot added the approved label Aug 2, 2024

google-oss-prow bot merged commit 00da4c9 into kubeflow:master Aug 2, 2024
39 checks passed

sandipanpanda deleted the jaxapi branch August 2, 2024 11:04

sandipanpanda mentioned this pull request Sep 19, 2024

[GSoC] Project 5: Integrate JAX with Kubeflow Training Operator to Support JAX Distributed Processes #2145

Open

sandipanpanda mentioned this pull request Oct 16, 2024

Add sandipanpanda as a Kubeflow Member kubeflow/internal-acls#718

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add JAX API #2163

Add JAX API #2163

sandipanpanda commented Jul 10, 2024

sandipanpanda commented Jul 10, 2024 •

edited

Loading

sandipanpanda commented Jul 10, 2024

coveralls commented Jul 10, 2024 •

edited

Loading

tenzen-y left a comment

tenzen-y Jul 12, 2024 •

edited

Loading

tenzen-y Jul 12, 2024 •

edited

Loading

andreyvelich left a comment

andreyvelich Jul 12, 2024

tenzen-y Jul 15, 2024

andreyvelich Jul 15, 2024 •

edited

Loading

tenzen-y Aug 1, 2024

sandipanpanda commented Jul 12, 2024

terrytangyuan left a comment

tenzen-y commented Aug 1, 2024

sandipanpanda commented Aug 1, 2024

andreyvelich commented Aug 1, 2024

tenzen-y commented Aug 2, 2024

google-oss-prow bot commented Aug 2, 2024

	func addJAXJobDefaultingFuncs(scheme *runtime.Scheme) error {
	func addJAXDefaultingFuncs(scheme *runtime.Scheme) error {

Add JAX API #2163

Add JAX API #2163

Conversation

sandipanpanda commented Jul 10, 2024

sandipanpanda commented Jul 10, 2024 • edited Loading

sandipanpanda commented Jul 10, 2024

coveralls commented Jul 10, 2024 • edited Loading

Pull Request Test Coverage Report for Build 9912856908

Warning: This coverage report may be inaccurate.

Details

💛 - Coveralls

tenzen-y left a comment

Choose a reason for hiding this comment

tenzen-y Jul 12, 2024 • edited Loading

Choose a reason for hiding this comment

tenzen-y Jul 12, 2024 • edited Loading

Choose a reason for hiding this comment

andreyvelich left a comment

Choose a reason for hiding this comment

andreyvelich Jul 12, 2024

Choose a reason for hiding this comment

tenzen-y Jul 15, 2024

Choose a reason for hiding this comment

andreyvelich Jul 15, 2024 • edited Loading

Choose a reason for hiding this comment

tenzen-y Aug 1, 2024

Choose a reason for hiding this comment

sandipanpanda commented Jul 12, 2024

terrytangyuan left a comment

Choose a reason for hiding this comment

tenzen-y commented Aug 1, 2024

sandipanpanda commented Aug 1, 2024

andreyvelich commented Aug 1, 2024

tenzen-y commented Aug 2, 2024

google-oss-prow bot commented Aug 2, 2024

sandipanpanda commented Jul 10, 2024 •

edited

Loading

coveralls commented Jul 10, 2024 •

edited

Loading

tenzen-y Jul 12, 2024 •

edited

Loading

tenzen-y Jul 12, 2024 •

edited

Loading

andreyvelich Jul 15, 2024 •

edited

Loading