ci(docker): modernize dependency management and Docker build pipeline

[garrett4wade]

Description

This PR overhauls the dependency management and Docker build infrastructure for improved reproducibility, faster builds, and better developer experience.

Key changes:

1. Dependency Management: Migrate to uv-based workflow with locked dependencies (uv.lock). Reorganize pyproject.toml with platform-specific extras (cuda, sglang, vllm, megatron) and dev dependencies. This ensures consistent environments across development, CI, and production.

2. Docker Build Optimization: Restructure Dockerfile with layered caching strategy - install torch and C++ packages (apex, grouped_gemm, transformer_engine) BEFORE syncing pyproject.toml dependencies. This prevents expensive recompilation when only Python dependencies change.

3. CI/CD Improvements:

   • Add runner-heartbeat workflow to prevent GitHub self-hosted runner expiration (7-day schedule)
   • Fix build-docker-image workflow to preserve running builder instances
   • Update tag-release-image to build fresh images instead of re-tagging
   • Update test-areal to use venv-based Python environment

4. Code Modernization:

   • Replace TypeVar with PEP 695 type parameter syntax (Python 3.12+)
   • Replace deprecated asyncio.TimeoutError with TimeoutError
   • Rename areal/tools/profile to profiling_utils to avoid import conflicts
   • Update vLLM imports for v0.14.0 compatibility

5. Documentation: Rewrite installation guide for uv-based workflow with clear instructions for CUDA/non-CUDA environments.

Impact:

• Faster Docker builds when only Python dependencies change (C++ packages stay cached)
• More reliable CI with self-hosted runner heartbeat mechanism
• Improved developer onboarding with clearer installation instructions
• Better reproducibility with locked dependencies

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update
• Code refactoring (no functional changes)
• Performance improvement
• Test coverage improvement

Checklist

• I have read the Contributing Guide
• I have run formatting tools (pre-commit or manual)
• I have run relevant unit tests and they pass
• I have added tests for new functionality
• I have updated documentation if needed
• My branch is up to date with main
• This PR introduces breaking changes (if yes, fill out details below)
• If this PR changes documentation, I have built and previewed it locally with jb build docs
• No critical issues raised by AI reviewers (/gemini review)

Breaking Change Details (if applicable):

This is a backward-compatible change. The new uv-based workflow is additive and does not remove existing pip-based installation options. Users can continue to use pip install -e . if they prefer.

Additional Context

File Change Summary

CI/CD Workflows:

• .github/workflows/build-docker-image.yml: Track whether builder was already running to avoid stopping it
• .github/workflows/runner-heartbeat.yml: New workflow to prevent runner expiration
• .github/workflows/tag-release-image.yml: Build fresh images on release instead of re-tagging
• .github/workflows/test-areal.yml: Use venv-based Python with explicit PATH setup

Docker:

• Dockerfile: Complete restructure with layered caching, upgrade to sglang 0.5.7 base image

Dependency Management:

• pyproject.toml: Reorganize with platform-specific extras and dev group
• uv.lock: New lockfile for reproducible builds (5000+ lines)

Documentation:

• docs/tutorial/installation.md: Rewrite for uv-based workflow
• evaluation/README.md, docs/tutorial/eval.md: Add evaluation dependency instructions

Code Modernization (no functional changes):

• Replace TypeVar with PEP 695 syntax in 5 files
• Replace asyncio.TimeoutError with TimeoutError in 10 files
• Rename areal/tools/profile to profiling_utils
• Update vLLM imports for v0.14.0

Testing

• Pre-commit hooks pass (Ruff, mdformat, nbstripout)
• Full integration tests require GPU cluster (not run locally)

[gemini-code-assist]

Summary of Changes

Hello @garrett4wade, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request modernizes the project's development and deployment environment by upgrading core deep learning dependencies and streamlining the package management workflow. It enhances build efficiency and reliability through updated Docker configurations and refined dependency installation practices, while also improving the internal validation system for installed packages.

Highlights

• Docker Image Update: The base Docker image has been updated from lmsysorg/sglang:v0.5.5.post1-cu129-amd64 to lmsysorg/sglang:v0.5.7-cu130-amd64-runtime, reflecting an upgrade to PyTorch 2.9.1 with CUDA 13.0 compatibility.
• Dependency Version Bumps: Key dependencies such as torch have been updated to 2.9.1+cu130, sglang to 0.5.7, and vllm to 0.14.0. Flash Attention installation now uses a specific pre-built wheel for 2.8.1+cu128torch2.9.
• Improved Dependency Management with uv: The Dockerfile, CONTRIBUTING.md, and README.md have been updated to leverage uv sync for more efficient and consistent dependency installation, including handling of optional extras and development dependencies.
• Parallelized Package Import Validation: The internal package validation logic in areal/tools/validation_base.py has been refactored to perform import tests in parallel using subprocesses, significantly speeding up dependency checks.
• Pre-commit Hook Adjustments: The .pre-commit-config.yaml now excludes uv.lock from large file checks and modifies Ruff linter/formatter to exclude evaluation and functioncall directories.
• Notebook Code Cleanup: Several Jupyter notebooks (math_reflection_en.ipynb, math_reflection_zh.ipynb, search_agent_zh.ipynb) have been cleaned up by removing unused imports and reformatting code for better readability.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request primarily focuses on updating the Docker image to a newer PyTorch version with CUDA 12.3.0 and refining dependency management using uv. The changes include updating base Docker images, adding necessary build tools like cmake, and restructuring uv dependency installation within the Dockerfile for better reproducibility and efficiency. Documentation and pre-commit hooks have been updated to reflect the new uv usage patterns. Additionally, the validation_base.py file has been refactored to improve the robustness and efficiency of package import tests by utilizing parallel subprocess execution. Overall, the changes are well-aligned with the stated objective of the pull request and represent good practices for dependency and build management.

[garrett4wade]

/gemini review

[gemini-code-assist]

Code Review

This is an excellent and comprehensive overhaul of the project's Docker build system and dependency management. The move to uv with staged builds, a virtual environment, and improved caching in the Dockerfile will significantly improve build times and reproducibility. The modernization of pyproject.toml to use dependency groups, platform markers for CUDA packages, and uv.lock is a fantastic step forward.

I've found one critical issue in the Dockerfile related to a hardcoded wheel path that will break the build, and a documentation error in one of the guides. Apart from these, the code modernization to Python 3.12 syntax and the refactoring in the tests are great improvements. Well done on this significant refactoring effort!

[rchardx]

LGTM. Waiting for CI.

[garrett4wade]

CI passes. Merging.