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]

This will change user workflow, can you revert it ?

[SurbhiJainUSC]

The reason for this change is that hf.login() opens an interactive shell whenever HF_TOKEN is unset. In non-interactive environments such as GKE, this causes the process to hang waiting for user input or fail.