Support JAX in Kueue using training-operator 1.9

[mimowo]

What would you like to be added:

Support for the recently added JAX in the training-operator, see kubeflow/training-operator#1619 and kubeflow/training-operator#1619.

Why is this needed:

For completeness of the support in Kueue. Also, the support for JAX was long-awaited in the training-operator, and there will be a community of people willing to use it with Kueue.

[mimowo]

cc @mbobrovskyi @mszadkow
Requires #4066

[tenzen-y]

I'm fine with supporting JAXJob in the Kueue!

Just FYI: we did not confirm if JAXJob is compatible with TPU 😞
Are you ok with such a JaxJob support level?

[mimowo]

IIUC, you haven't tested, but it does not indicate it is not working? I don't expect much of a difference between GPU and TPU, but without testing it remains unknown.

In any case I would be ok to claim the support for the GPUs.

[tenzen-y]

IIUC, you haven't tested, but it does not indicate it is not working? I don't expect much of a difference between GPU and TPU, but without testing it remains unknown.

In any case I would be ok to claim the support for the GPUs.

Actually, we did not have any verifications for both GPU and TPU. We verified only CPU

[mimowo]

I see, getting some confirmation that the controller works with GPU or TPU would be great.