Support DeepSWE evaluation pipeline on v5p-32 cluster with MaxText and Pathways

[susanbao]

Resolves #

Description

This PR adds support for the DeepSWE evaluation pipeline on v5p-32 cluster with MaxText and Pathways:

• Evaluation Pipeline (examples/deepswe/eval_deepswe.py):

  • Implements multi-slice evaluation with MaxText checkpoints and vLLM sampler on TPU v5p-32.
  • Adds direct MaxText checkpoint loading in vLLM to eliminate duplicate model allocation.
  • Supports direct GCS upload of evaluation results using google-cloud-storage Python client.
  • Adds memory cleanup (gc.collect(), del outputs) and continue_decode support.

• Resharding & Weight Transfer Optimization (tunix/generate/utils.py, tunix/rl/reshard.py, tunix/generate/vllm_sampler.py):

  • Adds zero-copy direct state update path when sharding and shape match between source and destination, bypassing resharding RPCs.
  • Optimizes chunked resharding in _reshard_in_chunks to update destination state in-place, deleting source chunk buffers and clearing caches to prevent memory accumulation in Pathways proxy.
  • Prioritizes jax.device_put for single-slice resharding and enables donate_input=True.

• Sampler & Trajectory Engine Support (tunix/generate/base_sampler.py, tunix/rl/agentic/trajectory/trajectory_collect_engine.py):

  • Supports both left_padded_prompt_tokens and padded_prompt_tokens on SamplerOutput.
  • Supports max_generation_steps and max_response_length forwarding in engine kwargs.

• TPU / Pathways Initialization & Checkpoint Support (scripts/patch_tpu_runner.py, tunix/models/automodel.py):

  • Adds scripts/patch_tpu_runner.py to patch rng_params_for_sampling directly on TPU mesh without nnx.Rngs / Pathways proxy compilation hangs.
  • Exposes checkpoint_storage_concurrent_gb in AutoModel for Orbax restore efficiency with large tensors.

• Dockerfile Updates:

  • Sets INSTALL_DEEPSWE_DEPS default to true, adds Dockerfile layer caching optimizations with CACHEBUST ARG, and sets full system PATH.

Checklist

• I have verified that my change does not break existing code and all unit tests pass.
• I have added all appropriate doc-strings/documentation.
• My PR is based on the latest changes of the main branch.
• I have signed the Contributor License Agreement https://cla.developers.google.com/about.
• I have followed Contribution Guidelines https://github.com/google/tunix/blob/main/docs/contributing.md.