Validate GPU support on JAX notebooks

[kaisoz]

GPU support on JAX notebooks couldn’t be validated due to an issue in our test infrastructure. To address this, work is ongoing to migrate the runners to CNCF’s GPU-based ARC in #3067.

Once that lands, we’ll need to test and validate the GPU support implementation for JAX notebooks and fix any bugs found during validation.

cc @kubeflow/kubeflow-trainer-team

[jaiakash]

Hi @kaisoz #3067 is merged, Newly added JAX is working correctly when i tried it, only thing is that since our CI has single GPU so i had to put the num_nodes to 1.

You can go forward with any testing for JAX notebooks.

[kaisoz]

Hi @kaisoz #3067 is merged, Newly added JAX is working correctly when i tried it, only thing is that since our CI has single GPU so i had to put the num_nodes to 1.

You can go forward with any testing for JAX notebooks.

Awesome! Thanks @jaiakash ! Do you know if there are any plans to have more GPUs in the future? Unfortunately we won't be able to test distributed training with this setup...

cc @andreyvelich

[jaiakash]

Yes we use Oracle Cloud based runner from CNCF. For this trainer, we are using GPU A10.1 single GPU, we do have GPU A10.2 with 2 GPUs available.
Ref: https://github.com/cncf/automation/tree/main/ci

I hope that would be sufficient for CI if CI is only use case right now.

But If you want to do any thing more specific, I can help to provision in cncfkubeflow tenancy.
Ref: https://docs.oracle.com/en-us/iaas/Content/Compute/References/computeshapes.htm#vm-gpu

Hope that helps.