Add JAX controller #2194

sandipanpanda · 2024-08-05T16:04:15Z

What this PR does / why we need it:
Implement JAX controller

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

/area gsoc

sandipanpanda · 2024-08-05T16:05:47Z

/cc @andreyvelich @tenzen-y @terrytangyuan

coveralls · 2024-08-05T16:08:53Z

Pull Request Test Coverage Report for Build 10960403998

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 10937611143:	0.0%
Covered Lines:	66
Relevant Lines:	66

💛 - Coveralls

pkg/controller.v1/jax/envvar.go

pkg/controller.v1/jax/jaxjob_controller.go

shravan-achar · 2024-08-14T22:53:58Z

@sandipanpanda - Is this PR ready for review?

tenzen-y · 2024-08-15T07:13:30Z

@sandipanpanda - Is this PR ready for review?

Actually, no. @sandipanpanda continues trying some implementations.
Regarding the remaining implementations, you can find the weekly sync documentation.

sandipanpanda · 2024-08-16T12:25:28Z

Adding jaxjob webhook test, examples and finishing up some work on test_e2e_jaxjob.py remain. Can you please share your input if the current implementation up until now is in the correct direction?

andreyvelich · 2024-08-16T14:06:32Z

@sandipanpanda If this PR is ready for review, please review the WIP from the PR title.

pkg/controller.v1/jax/jaxjob_controller_suite_test.go

sdk/python/test/e2e/test_e2e_jaxjob.py

sandipanpanda · 2024-08-16T19:42:38Z

I am still working on jaxjob webhook test, examples, example image for generate_container() in test_e2e_jaxjob.py. If that needs to be part of this PR, this is still WIP.

examples/jax/cpu-demo/train.py

examples/jax/cpu-demo/demo.yaml

pkg/controller.v1/jax/jaxjob_controller.go

tenzen-y · 2024-09-10T14:31:14Z

@sandipanpanda Could you address this error? We need to pass the appropriate arguments to e2e testing.

DEBUG kubernetes.client.rest:rest.py:235 response body: FATAL Flags parsing error:
flag --job_name=None: Flag --job_name must have a value other than None.
flag --sub_domain=None: Flag --sub_domain must have a value other than None.
flag --coordinator_port=None: Flag --coordinator_port must have a value other than None.
Pass --helpshort or --helpfull to see help on flags.

DEBUG kubernetes.client.rest:rest.py:235 response body: {"kind":"Status","apiVersion":"v1","metadata":{},"status":"Success","details":{"name":"jaxjob-cpu-ci-test","group":"kubeflow.org","kind":"jaxjobs","uid":"f7402a7c-588a-4b9d-903d-969ba0d4c7e2"}}

https://github.com/kubeflow/training-operator/actions/runs/10753814852/job/29823692795?pr=2194#step:9:5832

sandipanpanda · 2024-09-11T16:13:55Z

cc @tenzen-y PTAL

tenzen-y · 2024-09-11T16:28:07Z

examples/jax/cpu-demo/demo.yaml

+        spec:
+          containers:
+          - name: jax-worker
+            image: sandipanify/jaxgoogle


Suggested change

image: sandipanify/jaxgoogle

image: kubeflow/jaxjob-simple:latest

examples/jax/cpu-demo/Dockerfile

tenzen-y · 2024-09-11T16:31:07Z

examples/jax/cpu-demo/Dockerfile

+    libgoogle-glog-dev \
+    libgflags-dev \
+    libprotobuf-dev \
+    protobuf-compiler \


Library versions as well.

tenzen-y · 2024-09-11T16:31:37Z

pkg/controller.v1/jax/envvar.go

@sandipanpanda Missing unit testing.

pkg/controller.v1/jax/jaxjob_controller_test.go

tenzen-y · 2024-09-11T16:34:51Z

sdk/python/kubeflow/training/constants/constants.py

+JAXJOB_PLURAL = "jaxjobs"
+JAXJOB_CONTAINER = "jax"
+JAXJOB_REPLICA_TYPES = REPLICA_TYPE_WORKER.lower()
+JAXJOB_BASE_IMAGE = "docker.io/sandipanify/jaxgloo:latest"


Suggested change

JAXJOB_BASE_IMAGE = "docker.io/sandipanify/jaxgloo:latest"

JAXJOB_BASE_IMAGE = "kubeflow/jaxjob-simple:latest"

tenzen-y · 2024-09-11T16:35:56Z

sdk/python/test/e2e/test_e2e_jaxjob.py

+def generate_container() -> V1Container:
+    return V1Container(
+        name=CONTAINER_NAME,
+        image="docker.io/sandipanify/jaxgoogle:latest",


After this PR is merged and image is pushed into DockerHub, could you update this image in the followup PR?

Yes, I can do that.

Add JAX controller, controller tests, webhook validations, examples, e2e tests for JAXJob Extend the Training Operator Python SDK to simplify the creation and management of JAXJob resources. Signed-off-by: Sandipan Panda <samparksandipan@gmail.com>

pkg/controller.v1/jax/envvar_test.go

pkg/controller.v1/jax/envvar.go

pkg/controller.v1/jax/envvar_test.go

pkg/controller.v1/jax/envvar.go

Signed-off-by: Sandipan Panda <samparksandipan@gmail.com>

tenzen-y

Thanks for a lot of work!
/lgtm
/approve

google-oss-prow · 2024-09-20T14:33:52Z

[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

sandipanpanda · 2024-09-20T15:25:47Z

Thank you for your unwavering guidance and support throughout this project!

google-oss-prow bot added do-not-merge/work-in-progress area/gsoc labels Aug 5, 2024

google-oss-prow bot requested review from jinchihe and kuizhiqing August 5, 2024 16:04

google-oss-prow bot added the size/XL label Aug 5, 2024

google-oss-prow bot requested review from andreyvelich, tenzen-y and terrytangyuan August 5, 2024 16:05

terrytangyuan reviewed Aug 9, 2024

View reviewed changes

pkg/controller.v1/jax/envvar.go Show resolved Hide resolved

terrytangyuan reviewed Aug 9, 2024

View reviewed changes

pkg/controller.v1/jax/jaxjob_controller.go Show resolved Hide resolved

sandipanpanda force-pushed the jax-controller branch from cbb0ae0 to cf76911 Compare August 16, 2024 12:15

google-oss-prow bot added size/XXL and removed size/XL labels Aug 16, 2024

terrytangyuan reviewed Aug 16, 2024

View reviewed changes

pkg/controller.v1/jax/jaxjob_controller_suite_test.go Outdated Show resolved Hide resolved

sdk/python/test/e2e/test_e2e_jaxjob.py Outdated Show resolved Hide resolved

sandipanpanda force-pushed the jax-controller branch 3 times, most recently from a5e554b to 4c266aa Compare August 28, 2024 20:03

sandipanpanda changed the title ~~[WIP] Add JAX controller~~ Add JAX controller Aug 28, 2024

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

sandipanpanda force-pushed the jax-controller branch from 4c266aa to 212f8bf Compare August 28, 2024 20:07

tenzen-y reviewed Aug 30, 2024

View reviewed changes

sandipanpanda force-pushed the jax-controller branch from 212f8bf to 8e995a7 Compare September 4, 2024 19:13

sandipanpanda requested a review from tenzen-y September 4, 2024 19:14

sandipanpanda force-pushed the jax-controller branch 2 times, most recently from 28884e2 to 9c652b7 Compare September 7, 2024 19:14

sandipanpanda requested a review from tenzen-y September 7, 2024 19:34

sandipanpanda force-pushed the jax-controller branch 2 times, most recently from a3303b7 to f957d69 Compare September 10, 2024 20:17

tenzen-y reviewed Sep 11, 2024

View reviewed changes

sandipanpanda force-pushed the jax-controller branch 3 times, most recently from 9effb4a to 7c0563d Compare September 18, 2024 15:27

Add JAX controller

a86c927

Add JAX controller, controller tests, webhook validations, examples, e2e tests for JAXJob Extend the Training Operator Python SDK to simplify the creation and management of JAXJob resources. Signed-off-by: Sandipan Panda <samparksandipan@gmail.com>

sandipanpanda force-pushed the jax-controller branch from 7c0563d to 9219e94 Compare September 19, 2024 19:37

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 requested a review from tenzen-y September 19, 2024 20:42

tenzen-y reviewed Sep 19, 2024

View reviewed changes

sandipanpanda force-pushed the jax-controller branch from 9219e94 to 78623e1 Compare September 19, 2024 22:10

sandipanpanda requested a review from tenzen-y September 19, 2024 22:13

tenzen-y reviewed Sep 19, 2024

View reviewed changes

pkg/controller.v1/jax/envvar.go Outdated Show resolved Hide resolved

Add envvar tests

32aad29

Signed-off-by: Sandipan Panda <samparksandipan@gmail.com>

sandipanpanda force-pushed the jax-controller branch from 78623e1 to 32aad29 Compare September 20, 2024 13:43

sandipanpanda requested a review from tenzen-y September 20, 2024 13:43

tenzen-y reviewed Sep 20, 2024

View reviewed changes

google-oss-prow bot assigned tenzen-y Sep 20, 2024

google-oss-prow bot added the lgtm label Sep 20, 2024

google-oss-prow bot added the approved label Sep 20, 2024

google-oss-prow bot merged commit 8285aff into kubeflow:master Sep 20, 2024
40 checks passed

sandipanpanda deleted the jax-controller branch September 20, 2024 15:27

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add JAX controller #2194

Add JAX controller #2194

sandipanpanda commented Aug 5, 2024 •

edited

Loading

sandipanpanda commented Aug 5, 2024

coveralls commented Aug 5, 2024 •

edited

Loading

shravan-achar commented Aug 14, 2024 •

edited

Loading

tenzen-y commented Aug 15, 2024

sandipanpanda commented Aug 16, 2024

andreyvelich commented Aug 16, 2024

sandipanpanda commented Aug 16, 2024 •

edited

Loading

tenzen-y commented Sep 10, 2024

sandipanpanda commented Sep 11, 2024

tenzen-y Sep 11, 2024

tenzen-y Sep 11, 2024

tenzen-y Sep 11, 2024

tenzen-y Sep 11, 2024

tenzen-y Sep 11, 2024

sandipanpanda Sep 11, 2024

tenzen-y left a comment

google-oss-prow bot commented Sep 20, 2024

sandipanpanda commented Sep 20, 2024

	image: sandipanify/jaxgoogle
	image: kubeflow/jaxjob-simple:latest

	JAXJOB_BASE_IMAGE = "docker.io/sandipanify/jaxgloo:latest"
	JAXJOB_BASE_IMAGE = "kubeflow/jaxjob-simple:latest"

Add JAX controller #2194

Add JAX controller #2194

Conversation

sandipanpanda commented Aug 5, 2024 • edited Loading

sandipanpanda commented Aug 5, 2024

coveralls commented Aug 5, 2024 • edited Loading

Pull Request Test Coverage Report for Build 10960403998

Details

💛 - Coveralls

shravan-achar commented Aug 14, 2024 • edited Loading

tenzen-y commented Aug 15, 2024

sandipanpanda commented Aug 16, 2024

andreyvelich commented Aug 16, 2024

sandipanpanda commented Aug 16, 2024 • edited Loading

tenzen-y commented Sep 10, 2024

sandipanpanda commented Sep 11, 2024

tenzen-y Sep 11, 2024

Choose a reason for hiding this comment

tenzen-y Sep 11, 2024

Choose a reason for hiding this comment

tenzen-y Sep 11, 2024

Choose a reason for hiding this comment

tenzen-y Sep 11, 2024

Choose a reason for hiding this comment

tenzen-y Sep 11, 2024

Choose a reason for hiding this comment

sandipanpanda Sep 11, 2024

Choose a reason for hiding this comment

tenzen-y left a comment

Choose a reason for hiding this comment

google-oss-prow bot commented Sep 20, 2024

sandipanpanda commented Sep 20, 2024

sandipanpanda commented Aug 5, 2024 •

edited

Loading

coveralls commented Aug 5, 2024 •

edited

Loading

shravan-achar commented Aug 14, 2024 •

edited

Loading

sandipanpanda commented Aug 16, 2024 •

edited

Loading