Add Kubernetes Jobset support to `jax.distributed.initialize`

Open yhtang opened this issue 1 year ago • 1 comments

This should allow JAX to automatically bootstrap distributed computations when running inside a Kubernetes jobset.

Jul 08 '24 05:07 yhtang

@hawkinsp , PTAL when you can.

Sep 14 '24 13:09 mjsML

You need to add your new file to the list here: https://github.com/google/jax/blob/cd04d0f32e854aa754e37e4b676725655a94e731/jax/BUILD#L949

(sorted order)

to unbreak the bazel build.

added

Sep 20 '24 17:09 yhtang