Testing CI in JAX example (kubeflow/trainer#2385)

* Add MNIST example with SPMD for JAX

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

Use -- server-side to install the latest local changes of Training
Operator control plane

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

* Add JAXJob output

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

* Update JAXJob CI images

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

* Adjust jaxjob spmd example batch size

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

* Add JAX Example Docker Image Build in CI

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

* Fix script name typo

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

* Update script permissions

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

* Add KIND_CLUSTER env var

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

* Increase timeouts

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

* Test higher resources

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

* Increase Timeout

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

* remove resource reqs

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

* test low batch size

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

* test small batch size

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

* Hardcode number of batches

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

---------

Signed-off-by: Sandipan Panda <[email protected]>
Signed-off-by: sailesh duddupudi <[email protected]>
Co-authored-by: Sandipan Panda <[email protected]>

Author: saileshd1402

python/kubeflow/training/constants/constants.py

@@ -153,7 +153,7 @@
 JAXJOB_PLURAL = "jaxjobs"
 JAXJOB_CONTAINER = "jax"
 JAXJOB_REPLICA_TYPES = REPLICA_TYPE_WORKER.lower()
-JAXJOB_BASE_IMAGE = "docker.io/kubeflow/jaxjob-simple:latest"
+JAXJOB_BASE_IMAGE = "docker.io/kubeflow/jaxjob-dist-spmd-mnist:latest"
 
 # Dictionary to get plural, model, and container for each Job kind.
 JOB_PARAMETERS = {

python/test/e2e/test_e2e_jaxjob.py

@@ -155,7 +155,6 @@ def generate_jaxjob(
 def generate_container() -> V1Container:
     return V1Container(
         name=CONTAINER_NAME,
-        image="docker.io/kubeflow/jaxjob-simple:latest",
-        command=["python", "train.py"],
-        resources=V1ResourceRequirements(limits={"memory": "2Gi", "cpu": "0.8"}),
+        image=os.getenv("JAX_JOB_IMAGE", "docker.io/kubeflow/jaxjob-dist-spmd-mnist:latest"),
+        resources=V1ResourceRequirements(limits={"memory": "3Gi", "cpu": "1.2"}),
     )