Testing CI in JAX example #2385

saileshd1402 · 2025-01-09T19:27:03Z

What this PR does / why we need it:

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

Illustrate how to use JAX's `pmap` to express and execute single-program multiple-data (SPMD) programs for data parallelism along a batch dimension Signed-off-by: Sandipan Panda <[email protected]>

Use -- server-side to install the latest local changes of Training Operator control plane Signed-off-by: Sandipan Panda <[email protected]>

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

Signed-off-by: sailesh duddupudi <[email protected]>

coveralls · 2025-01-09T19:32:13Z

Pull Request Test Coverage Report for Build 12815665472

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 12718184264:	0.0%
Covered Lines:	85
Relevant Lines:	85

💛 - Coveralls

Signed-off-by: sailesh duddupudi <[email protected]>

andreyvelich · 2025-01-15T22:47:45Z

Hi @saileshd1402 @sandipanpanda @mahdikhashan, any updates with the JAX example ?
Since we want to release 1.9 soon, do we want to include this PR as part of the release ?

andreyvelich · 2025-01-15T22:47:57Z

cc @Electronic-Waste @kubeflow/wg-training-leads

Electronic-Waste · 2025-01-16T00:11:23Z

@andreyvelich It would be great if we could catch up with the final feature freeze day on Jan 20th. However, we could implement the CI for JAX example with notebooks like #2246 if the integration test cannot be ready before Jan 20th.

saileshd1402 · 2025-01-16T03:08:34Z

I'll try to make the job use lesser resources by the end of today. But if it's not working, then I agree with @Electronic-Waste

mahdikhashan · 2025-01-16T09:03:20Z

Hi @saileshd1402 @sandipanpanda @mahdikhashan, any updates with the JAX example ? Since we want to release 1.9 soon, do we want to include this PR as part of the release ?

hi, i couldn't figure out what could be changed, not yet, regarding the first epoch. i'll spend more time tomorrow.

saileshd1402 · 2025-01-16T11:11:21Z

examples/jax/jax-dist-spmd-mnist/spmd_mnist_classifier_fromscratch.py

+
+    for epoch in range(num_epochs):
+        start_time = time.time()
+        num_batches = 5


@andreyvelich @mahdikhashan The tests are working when I just hard coded this num_batches test to check CI, it's passing so if we make some params smaller specific to CI, the tests should pass. Is there anywhere we can compromise for CI?

We can keep this number of batches small.
Was it the main issue why the first epoch takes ~ 1 hour ?

Yes, from my understanding this was the reason. Resources was not enough to handle more batches/takes way too long.

Should I edit the script so that if a certain CI env var is set, then we set num_batches to lower value?

Let's use this number of batches by default.
You can just add the comment that increasing number of batches requires more resources.

Signed-off-by: sailesh duddupudi <[email protected]>

saileshd1402 · 2025-01-16T19:39:48Z

@andreyvelich I have addressed your comment. Seems like CI looks good now, please review when you're available.

cc: @sandipanpanda @mahdikhashan @Electronic-Waste

andreyvelich

Thank you for this effort @saileshd1402!
/lgtm
/assign @Electronic-Waste @kubeflow/wg-training-leads @sandipanpanda

Electronic-Waste

LGTM. Thanks for your great contribution! @saileshd1402

/lgtm

andreyvelich

Thanks @saileshd1402!
/approve

google-oss-prow · 2025-01-17T17:07:35Z

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: andreyvelich, Electronic-Waste

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~~ [andreyvelich]

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

andreyvelich · 2025-01-17T17:42:05Z

/cherry-pick release-1.9

google-oss-robot · 2025-01-17T17:42:56Z

@andreyvelich: new pull request created: #2390

In response to this:

/cherry-pick release-1.9

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/test-infra repository.

sandipanpanda and others added 6 commits December 26, 2024 01:30

Add MNIST example with SPMD for JAX

619aa20

Illustrate how to use JAX's `pmap` to express and execute single-program multiple-data (SPMD) programs for data parallelism along a batch dimension Signed-off-by: Sandipan Panda <[email protected]>

Update CONTRIBUTING.md

ca78118

Use -- server-side to install the latest local changes of Training Operator control plane Signed-off-by: Sandipan Panda <[email protected]>

Add JAXJob output

7f523e4

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

Update JAXJob CI images

1ed9b7d

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

Adjust jaxjob spmd example batch size

904dbd6

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

Add JAX Example Docker Image Build in CI

ef997a1

Signed-off-by: sailesh duddupudi <[email protected]>

google-oss-prow bot added the do-not-merge/work-in-progress label Jan 9, 2025

google-oss-prow bot requested review from jinchihe and kuizhiqing January 9, 2025 19:27

google-oss-prow bot added the size/L label Jan 9, 2025

saileshd1402 added 4 commits January 9, 2025 19:36

Fix script name typo

9a88072

Signed-off-by: sailesh duddupudi <[email protected]>

Update script permissions

a2ca3a3

Signed-off-by: sailesh duddupudi <[email protected]>

Add KIND_CLUSTER env var

15e6205

Signed-off-by: sailesh duddupudi <[email protected]>

Increase timeouts

07688ec

Signed-off-by: sailesh duddupudi <[email protected]>

google-oss-prow bot added size/XL and removed size/L labels Jan 9, 2025

saileshd1402 commented Jan 16, 2025

View reviewed changes

saileshd1402 added 6 commits January 16, 2025 18:50

Test higher resources

b876b7b

Signed-off-by: sailesh duddupudi <[email protected]>

Increase Timeout

063adfa

Signed-off-by: sailesh duddupudi <[email protected]>

remove resource reqs

76f2e80

Signed-off-by: sailesh duddupudi <[email protected]>

test low batch size

3353f7a

Signed-off-by: sailesh duddupudi <[email protected]>

test small batch size

c80db23

Signed-off-by: sailesh duddupudi <[email protected]>

Hardcode number of batches

343090c

Signed-off-by: sailesh duddupudi <[email protected]>

saileshd1402 force-pushed the jax-example branch from bb9222a to 343090c Compare January 16, 2025 18:54

saileshd1402 changed the title ~~[WIP] Testing CI in JAX example~~ Testing CI in JAX example Jan 16, 2025

google-oss-prow bot removed the do-not-merge/work-in-progress label Jan 16, 2025

andreyvelich reviewed Jan 16, 2025

View reviewed changes

google-oss-prow bot assigned andreyvelich Jan 16, 2025

google-oss-prow bot added the lgtm label Jan 16, 2025

Electronic-Waste approved these changes Jan 17, 2025

View reviewed changes

google-oss-prow bot assigned Electronic-Waste Jan 17, 2025

andreyvelich reviewed Jan 17, 2025

View reviewed changes

google-oss-prow bot added the approved label Jan 17, 2025

google-oss-prow bot merged commit 6d58ea9 into kubeflow:master Jan 17, 2025
56 checks passed

google-oss-robot mentioned this pull request Jan 17, 2025

[release-1.9] Testing CI in JAX example #2390

Merged

andreyvelich mentioned this pull request Jan 17, 2025

Track follow-up for adding CI image build to the Kind cluster #2364

Closed

1 task

sandipanpanda mentioned this pull request Mar 27, 2025

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

Closed

Testing CI in JAX example #2385

Testing CI in JAX example #2385

Uh oh!

Conversation

saileshd1402 commented Jan 9, 2025

Uh oh!

coveralls commented Jan 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Pull Request Test Coverage Report for Build 12815665472

Details

💛 - Coveralls

Uh oh!

andreyvelich commented Jan 15, 2025

Uh oh!

andreyvelich commented Jan 15, 2025

Uh oh!

Electronic-Waste commented Jan 16, 2025

Uh oh!

saileshd1402 commented Jan 16, 2025

Uh oh!

mahdikhashan commented Jan 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

saileshd1402 Jan 16, 2025

Choose a reason for hiding this comment

Uh oh!

andreyvelich Jan 16, 2025

Choose a reason for hiding this comment

Uh oh!

saileshd1402 Jan 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

andreyvelich Jan 16, 2025

Choose a reason for hiding this comment

Uh oh!

saileshd1402 commented Jan 16, 2025

Uh oh!

andreyvelich left a comment

Choose a reason for hiding this comment

Uh oh!

Electronic-Waste left a comment

Choose a reason for hiding this comment

Uh oh!

andreyvelich left a comment

Choose a reason for hiding this comment

Uh oh!

google-oss-prow bot commented Jan 17, 2025

Uh oh!

Uh oh!

andreyvelich commented Jan 17, 2025

Uh oh!

google-oss-robot commented Jan 17, 2025

Uh oh!

Uh oh!

coveralls commented Jan 9, 2025 •

edited

Loading

mahdikhashan commented Jan 16, 2025 •

edited

Loading

saileshd1402 Jan 16, 2025 •

edited

Loading