Add MaxText and Raiden optional dependencies to Dockerfile

[SurbhiJainUSC]

Summary

This PR adds modular build support for MaxText and Raiden dependencies in Dockerfile via conditional build arguments (INSTALL_MAXTEXT and INSTALL_RAIDEN), allowing reproducible container builds for distributed RL workloads without bloating or breaking the standard base image.

Key Changes

1. Dockerfile:
   • Installs uv directly in the virtual environment (/opt/venv).
   • Adds ARG INSTALL_MAXTEXT=false to conditionally install MaxText/Flax training dependencies from requirements/maxtext_requirements.txt.
   • Adds ARG INSTALL_RAIDEN=false to conditionally install Raiden weight synchronization wheels (/app/raiden_wheels/*.whl) or tpu-raiden-jax from Google Artifact Registry.
2. requirements/maxtext_requirements.txt (maxtext_requirements.txt):
   • Defines standalone dependencies required for running MaxText trainer backends (e.g. maxtext, flax, optax).
3. requirements/special_requirements.txt (special_requirements.txt):
   • Pins compatible package versions for distributed RL runtimes.
4. tunix/oss/utils.py (utils.py):
   • Refines OSS utility helper imports and formatting.

Checklist

• I have added all the necessary unit tests for my change.
• 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 (if unsure, rebase the code).
• I have signed the Contributor License Agreement.
• I have followed Contribution Guidelines.

[jiangyangmu]

Bump this version should be careful, we need to run package test and nightly regression test before bump.

Do you depend on changes in the updated version ?

[SurbhiJainUSC]

Yes, we need these weight sync changes in tpu-inference: vllm-project/tpu-inference@4a89516

[wang2yn84]

When you update tpu-inference, please also update vllm, and the vllm nightly image tag in https://github.com/google/tunix/tree/main/.github/workflows. So the correct way of doing this is identify which vllm docker image you'd like to use, then apply the same commit to these requirement files.