[tx] replace calls to jax.process_index() to resolve rank ordering issue with multi-host TPUs

[andrewsykim]

jax.process_index() is inconsistent between GPU and TPU backends. For GPU / CPU, it returns the process_id passed into jax.distributed.initialize. However, for TPU, the index assignment happens from the TPU slice, regardless of what is assigned from jax.distributed.initialize or the TPU environment (e.g. TPU_WORKER_ID from GKE). This leads tojax.process_index() returning a different rank from --process-id set from the user.

Assuming jax.process_index() won't be updated in the near future to be consistent across TPU / GPU, this PR updates the SkyRL JAX backend to only use the index provided from the --process-id argument.

I've tested this PR on my 4x4 TPU v6e cluster with both Qwen3-0.6B and Qwen3-8B.

Resolves #1024

---

[pcmoritz]

We will also need to update

SkyRL/skyrl-tx/tx/utils/models.py

Line 300 in 675851c

with pack_and_upload(output_path, rank=jax.process_index()) as temp_dir:

so the checkpointing works correctly in the multi-node setting :)

[andrewsykim]

so the checkpointing works correctly

Updated

[pcmoritz]

It would probably be better to remove the default from this, otherwise it is too easy to make a mistake and not specify this in call sites (before, it was kind of doing implicitly the right thing by using jax.process_index(), but now it is important to pass in the right value here).

I can make that change before merging the PR.

[pcmoritz]

/gemini review