Add JAX controller

[sandipanpanda]

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):
Fix: #1619
Ref:#2145

• https://jax.readthedocs.io/en/latest/multi_process.html
• https://jax.readthedocs.io/en/latest/_autosummary/jax.distributed.initialize.html#jax.distributed.initialize

/area gsoc

[sandipanpanda]

/cc @andreyvelich @tenzen-y @terrytangyuan

[coveralls]

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

[shravan-achar]

@sandipanpanda - Is this PR ready for review?

[tenzen-y]

@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]

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]

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

[tenzen-y]

@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]

cc @tenzen-y PTAL

[tenzen-y]

Thanks for a lot of work!
/lgtm
/approve

[google-oss-prow]

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

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