JAX Integration Enhancement Proposal #2125

sandipanpanda · 2024-05-22T19:46:42Z

Enhancement Proposal for integrating JAX with Kubeflow Training Operator for distributed training on Kubernetes

Ref:

/area gsoc

sandipanpanda · 2024-05-23T11:27:45Z

/cc @andreyvelich @tenzen-y @terrytangyuan @johnugeorge

kannon92 · 2024-05-24T15:25:19Z

@andreyvelich @tenzen-y do we want to see about using JobSet here? API seems pretty simple so far and JobSet should be a quick drop in from what I see.

andreyvelich · 2024-05-24T19:26:09Z

@andreyvelich @tenzen-y do we want to see about using JobSet here? API seems pretty simple so far and JobSet should be a quick drop in from what I see.

Yes, we will work with @sandipanpanda to refactor this proposal based on our discussion from this doc: https://docs.google.com/document/d/1bha8rB6_iPTi9rXjJMnKi-CLxfL7dwtmQCKH4N6dcsI/edit#heading=h.24ap79rbqoxl.
E.g. integrate Jax to Training Operator with JobSet and TrainJob.

coveralls · 2024-05-24T19:29:27Z

Pull Request Test Coverage Report for Build 9207133318

Details

0 of 0 changed or added relevant lines in 0 files are covered.
No unchanged relevant lines lost coverage.
Overall coverage increased (+0.02%) to 35.431%

Totals
Change from base Build 9179239826:	0.02%
Covered Lines:	4380
Relevant Lines:	12362

💛 - Coveralls

docs/proposals/jax-integration.md

tenzen-y · 2024-06-03T15:59:10Z

docs/proposals/jax-integration.md

+  - **Mitigation**: Implement robust scheduling and resource management policies in the Training Operator.
+
+## Design Details
+


Could you add a corresponding table between all JAX environment variables and how to obtain OR configure the JAX parameters like this?

https://docs.google.com/document/d/1bha8rB6_iPTi9rXjJMnKi-CLxfL7dwtmQCKH4N6dcsI/edit#heading=h.d6pih1ui0ccr

@sandipanpanda Is there any tables?

docs/proposals/jax-integration.md

terrytangyuan · 2024-06-11T13:30:17Z

docs/proposals/jax-integration.md

+  selector:
+    app: jaxjob-master-${job_id}
+  ports:
+  - port: 23456


Document the default port somewhere

docs/proposals/jax-integration.md

andreyvelich · 2024-06-18T19:10:16Z

@sandipanpanda We agreed that initially we can support Jax within existing Training Operator implementation given the time it takes for us to design and to implement TrainJob API with JobSet.
Once we implement TrainJob API in Training Operator, we will easily add support for Jax settings.
We are planning to implement TrainJob API in Q4 - Q1.

That will unblock @sandipanpanda with this work.
@sandipanpanda Feel free to finalize your GSoC proposal using the existing Training Operator assets.
cc @kubeflow/wg-training-leads

sandipanpanda · 2024-06-19T20:11:04Z

Thanks! I have updated the proposal according to Kubeflow Proposal Document Structure. I am working on finalizing the proposal using the existing Training Operator assets.

sandipanpanda · 2024-07-09T15:33:04Z

/cc @andreyvelich @tenzen-y PTAL

google-oss-prow · 2024-07-09T15:33:08Z

@sandipanpanda: GitHub didn't allow me to request PR reviews from the following users: PTAL.

Note that only kubeflow members and repo collaborators can review this PR, and authors cannot review their own PRs.

In response to this:

/cc @andreyvelich @tenzen-y PTAL

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository.

tenzen-y

Additionally, could you mention about validations?
You can refer to the other framework validations like this: https://github.com/kubeflow/training-operator/blob/2b39d3cbc1525dc29eb40aa58b149dbeef00aa0f/pkg/webhooks/pytorch/pytorchjob_webhook.go

docs/proposals/jax-integration.md

tenzen-y · 2024-07-10T09:28:09Z

docs/proposals/jax-integration.md

+  - **Mitigation**: Implement robust scheduling and resource management policies in the Training Operator.
+
+## Design Details
+


@sandipanpanda Is there any tables?

sandipanpanda · 2024-07-11T17:47:36Z

@tenzen-y PTAL

coveralls · 2024-07-11T18:29:27Z

Pull Request Test Coverage Report for Build 9898700147

Details

0 of 0 changed or added relevant lines in 0 files are covered.
No unchanged relevant lines lost coverage.
Overall coverage increased (+0.02%) to 35.423%

Totals
Change from base Build 9896165546:	0.02%
Covered Lines:	4380
Relevant Lines:	12365

💛 - Coveralls

tenzen-y

Otherwise lgtm

docs/proposals/jax-integration.md

tenzen-y · 2024-07-11T18:47:47Z

docs/proposals/jax-integration.md

+```yaml
+apiVersion: kubeflow.org/v1
+kind: JAXJob
+metadata:
+  name: example-jaxjob
+spec:
+  backend: "gloo"
+  coordinatorPort: "6666"
+  replicaSpecs:
+    - replicas: 1
+      ReplicaType: COORDINATOR
+      template:
+      spec:
+        containers:
+          - image: ghcr.io/kubeflow/jax:latest
+            name: coordinator
+            imagePullPolicy: IfNotPresent
+        restartPolicy: OnFailure
+    - replicas: 2
+      ReplicaType: WORKER
+      template:
+        spec:
+          containers:
+            - image: ghcr.io/kubeflow/jax:latest
+              name: worker
+          restartPolicy: OnFailure
+```


This YAML example is abandoned. So, could you refine this based on the latest API design?

tenzen-y

Actually, I never find validation design through this design.
Could you add which validations should we add?
For example, do we need to validate if the JaxJob has cordinator role?

Kubeflow Enhancement Proposal: Integrate JAX with Kubeflow Training Operator Signed-off-by: Sandipan Panda <samparksandipan@gmail.com>

tenzen-y

Thank you!
/lgtm
/approve

google-oss-prow · 2024-07-12T10:27:06Z

[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 area/gsoc label May 22, 2024

google-oss-prow bot requested review from jinchihe and kuizhiqing May 22, 2024 19:46

google-oss-prow bot added the size/L label May 22, 2024

sandipanpanda force-pushed the jax-integration-proposal branch from c45d88e to 7ca6f7e Compare May 23, 2024 11:25

google-oss-prow bot requested review from andreyvelich, johnugeorge, tenzen-y and terrytangyuan May 23, 2024 11:27

tenzen-y reviewed May 31, 2024

View reviewed changes

docs/proposals/jax-integration.md Outdated Show resolved Hide resolved

tenzen-y reviewed Jun 3, 2024

View reviewed changes

terrytangyuan reviewed Jun 11, 2024

View reviewed changes

docs/proposals/jax-integration.md Outdated Show resolved Hide resolved

sandipanpanda mentioned this pull request Jun 13, 2024

[GSOC] Tracking Issue: Integrate JAX in Kubeflow Training Operator #2145

Open

12 tasks

shravan-achar reviewed Jun 18, 2024

View reviewed changes

docs/proposals/jax-integration.md Outdated Show resolved Hide resolved

docs/proposals/jax-integration.md Outdated Show resolved Hide resolved

docs/proposals/jax-integration.md Outdated Show resolved Hide resolved

sandipanpanda force-pushed the jax-integration-proposal branch from 7ca6f7e to 2040ac4 Compare June 19, 2024 20:03

sandipanpanda force-pushed the jax-integration-proposal branch from 2040ac4 to abc4ca4 Compare June 28, 2024 14:39

sandipanpanda force-pushed the jax-integration-proposal branch from abc4ca4 to 1e8059d Compare July 5, 2024 14:53

google-oss-prow bot added size/XL and removed size/L labels Jul 5, 2024

sandipanpanda force-pushed the jax-integration-proposal branch from 1e8059d to 895620d Compare July 8, 2024 20:20

google-oss-prow bot added size/L and removed size/XL labels Jul 8, 2024

sandipanpanda force-pushed the jax-integration-proposal branch from 895620d to 242df25 Compare July 8, 2024 20:27

google-oss-prow bot requested a review from tenzen-y July 9, 2024 15:33

tenzen-y reviewed Jul 10, 2024

View reviewed changes

sandipanpanda force-pushed the jax-integration-proposal branch 3 times, most recently from 9c4c9bd to e3eaa99 Compare July 11, 2024 17:46

tenzen-y reviewed Jul 11, 2024

View reviewed changes

JAX Integration Enhancement Proposal

816efcc

Kubeflow Enhancement Proposal: Integrate JAX with Kubeflow Training Operator Signed-off-by: Sandipan Panda <samparksandipan@gmail.com>

sandipanpanda force-pushed the jax-integration-proposal branch from e3eaa99 to 816efcc Compare July 11, 2024 21:24

tenzen-y reviewed Jul 12, 2024

View reviewed changes

google-oss-prow bot assigned tenzen-y Jul 12, 2024

google-oss-prow bot added the lgtm label Jul 12, 2024

google-oss-prow bot added the approved label Jul 12, 2024

google-oss-prow bot merged commit bcba864 into kubeflow:master Jul 12, 2024
39 checks passed

sandipanpanda deleted the jax-integration-proposal branch July 12, 2024 19:52

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

JAX Integration Enhancement Proposal #2125

JAX Integration Enhancement Proposal #2125

sandipanpanda commented May 22, 2024

sandipanpanda commented May 23, 2024

kannon92 commented May 24, 2024

andreyvelich commented May 24, 2024

coveralls commented May 24, 2024

tenzen-y Jun 3, 2024

tenzen-y Jul 10, 2024

terrytangyuan Jun 11, 2024

andreyvelich commented Jun 18, 2024

sandipanpanda commented Jun 19, 2024

sandipanpanda commented Jul 9, 2024

google-oss-prow bot commented Jul 9, 2024

tenzen-y left a comment

tenzen-y Jul 10, 2024

sandipanpanda commented Jul 11, 2024

coveralls commented Jul 11, 2024 •

edited

Loading

tenzen-y left a comment

tenzen-y Jul 11, 2024

tenzen-y left a comment

tenzen-y left a comment

google-oss-prow bot commented Jul 12, 2024

		- Mitigation: Implement robust scheduling and resource management policies in the Training Operator.

		## Design Details

JAX Integration Enhancement Proposal #2125

JAX Integration Enhancement Proposal #2125

Conversation

sandipanpanda commented May 22, 2024

sandipanpanda commented May 23, 2024

kannon92 commented May 24, 2024

andreyvelich commented May 24, 2024

coveralls commented May 24, 2024

Pull Request Test Coverage Report for Build 9207133318

Details

💛 - Coveralls

tenzen-y Jun 3, 2024

Choose a reason for hiding this comment

tenzen-y Jul 10, 2024

Choose a reason for hiding this comment

terrytangyuan Jun 11, 2024

Choose a reason for hiding this comment

andreyvelich commented Jun 18, 2024

sandipanpanda commented Jun 19, 2024

sandipanpanda commented Jul 9, 2024

google-oss-prow bot commented Jul 9, 2024

tenzen-y left a comment

Choose a reason for hiding this comment

tenzen-y Jul 10, 2024

Choose a reason for hiding this comment

sandipanpanda commented Jul 11, 2024

coveralls commented Jul 11, 2024 • edited Loading

Pull Request Test Coverage Report for Build 9898700147

Details

💛 - Coveralls

tenzen-y left a comment

Choose a reason for hiding this comment

tenzen-y Jul 11, 2024

Choose a reason for hiding this comment

tenzen-y left a comment

Choose a reason for hiding this comment

tenzen-y left a comment

Choose a reason for hiding this comment

google-oss-prow bot commented Jul 12, 2024

coveralls commented Jul 11, 2024 •

edited

Loading