docs: add checkpointing, metrics tracking, and workflow guides

[garrett4wade]

Description

Add comprehensive reference documentation for core AReaL systems including checkpointing,
metrics tracking, rollout workflows, and allocation modes. Also includes recovery system
refactoring with simplified mode handling and unit tests.

Related Issue

N/A

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):

N/A

Additional Context

Key changes:

New Documentation:

• docs/reference/checkpointing.md: Saver (HF format) and RecoverHandler (DCP) docs
• docs/reference/metrics_tracking.md: Streaming vs batch paradigms, distributed aggregation
• docs/reference/rollout_workflow.md: Interface, return types, implementation patterns
• docs/reference/alloc_mode.md: Resource allocation strategies reference
• docs/best_practices/workflow.md: Workflow implementation best practices

Code Changes:

• areal/utils/recover.py: Simplified mode handling (deprecate fault/resume in favor of on/off/auto/disabled)
• areal/api/cli_args.py: Add mode validation with migration hints
• areal/tests/test_recover.py: Unit tests for recovery configuration

Tutorial Reorganization:

• Removed docs/tutorial/troubleshooting.md (content moved to reference docs)
• Expanded docs/tutorial/eval.md with more details
• Clarified docs/tutorial/gsm8k_grpo.md architecture walkthrough

[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 significantly enhances the AReaL project's documentation by introducing detailed guides and references for critical features. It covers checkpointing for both model export and fault recovery, a unified metrics tracking system for distributed training, and a comprehensive explanation of GPU allocation modes. Furthermore, it provides an in-depth reference for the RolloutWorkflow API and best practices for developing agent workflows. The changes also include a structural reorganization of the documentation and the removal of outdated content, ensuring that users have access to up-to-date and clear information for developing and managing AReaL experiments.

Highlights

• Comprehensive Documentation Added: Introduced new, in-depth documentation covering AReaL's checkpointing system, metrics tracking (including WandB and TensorBoard integration), GPU allocation modes, and the RolloutWorkflow API reference.
• Best Practices for Agent Workflows: A new guide on best practices for writing efficient agent workflows has been added, emphasizing asynchronous operations, proper reward function wrapping, and HTTP client management.
• Documentation Structure Reorganization: The overall documentation structure has been reorganized, including updates to _toc.yml, to improve navigability and integrate the new reference sections seamlessly.
• Removal of Deprecated Features: The unused AREAL_RECOVER_RUN environment variable has been removed from several core Python files, and the deprecated troubleshooting guide has been eliminated from the documentation.
• Enhanced Tutorial Clarity: Existing tutorials, such as those for evaluation and Megatron integration, have been expanded and clarified to align with the new documentation and best practices.

🧠 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.

Changelog

• areal/launcher/local.py
  • Removed AREAL_RECOVER_RUN environment variable.
• areal/launcher/ray.py
  • Removed AREAL_RECOVER_RUN environment variable.
• areal/launcher/slurm.py
  • Removed AREAL_RECOVER_RUN environment variable.
• areal/scheduler/slurm.py
  • Removed os import due to the removal of os.getenv for AREAL_RECOVER_RUN.
  • Removed AREAL_RECOVER_RUN environment variable forwarding.
• areal/utils/recover.py
  • Removed the check for AREAL_RECOVER_RUN environment variable in the load method.
• docs/_toc.yml
  • Restructured the table of contents to include new 'Reference' and 'Code Walkthrough' sections.
  • Added new documentation files: reference/checkpointing.md, reference/metrics_tracking.md, reference/alloc_mode.md, reference/rollout_workflow.md, and best_practices/workflow.md.
  • Removed tutorial/openai_workflows.
  • Moved tutorial/gsm8k_grpo to the new 'Code Walkthrough' section.
• docs/best_practices/workflow.md
  • Added a new document detailing best practices for writing agent workflows, covering async/await, AsyncRewardWrapper, initialization, and HTTP client reuse.
• docs/reference/alloc_mode.md
  • Added a new document explaining AReaL's GPU allocation mode system, including syntax, parallelism dimensions, and backend selection.
• docs/reference/checkpointing.md
  • Added a new document describing AReaL's checkpointing mechanisms (Saver for HuggingFace export and RecoverHandler for fault tolerance).
• docs/reference/metrics_tracking.md
  • Added a new document outlining AReaL's unified metrics tracking system, including streaming/batch metrics, API, and logging backends.
• docs/reference/rollout_workflow.md
  • Added a new document providing a reference for the RolloutWorkflow abstraction, its interface, and implementation details.
• docs/tutorial/agentic_rl.md
  • Updated a link to the workflow/ directory to point to the GitHub repository.
• docs/tutorial/eval.md
  • Expanded the introduction to clarify AReaL's evaluation scope and compatibility with third-party frameworks.
  • Added a new section on 'Evaluation Metrics' with a code example.
  • Updated the architecture description and code snippet for evaluation setup, making scheduler and engine initialization more explicit.
  • Revised the logging section from 'Logging Metrics' to 'Logging Results' and simplified the logging example.
  • Updated internal links to use ref for distributed experiments.
• docs/tutorial/gsm8k_grpo.md
  • Adjusted ASCII art diagrams for 'Data Flow with RTensor' and 'Three-Process Architecture' for better alignment.
  • Clarified that rollout and training happen simultaneously.
  • Removed the 'Overlap with Training' and 'Pause/Resume' sections, as the simultaneous nature is now stated earlier.
  • Updated the description of NCCL-based transfer for weight synchronization.
  • Clarified the weight sync process steps.
  • Expanded the 'Checkpointing' section to detail both Saver and RecoverHandler mechanisms and linked to the new Checkpointing Reference.
  • Expanded the 'Metrics Tracking' section to describe streaming and batch metrics, provide code examples, and link to the new Metrics Tracking Reference.
  • Added 'Evaluation' to 'Next Steps'.
• docs/tutorial/megatron.md
  • Simplified the introduction by removing redundant explanations about backend auto-selection.
  • Replaced the 'Understanding allocation_mode' section with a link to the new 'Allocation Mode Reference'.
  • Refactored 'Training Parallel Strategy' and 'Inference Parallel Strategy' into a more concise 'MoE Parallel Strategy' section, linking to the new 'Allocation Mode Reference' for details.
• docs/tutorial/troubleshooting.md
  • Removed as deprecated.

Activity

• The pull request was authored by garrett4wade, focusing on a significant documentation overhaul.
• New documentation files were added for checkpointing, metrics tracking, GPU allocation modes, and the RolloutWorkflow API.
• Existing documentation was updated and reorganized for improved clarity and structure.
• An unused environment variable, AREAL_RECOVER_RUN, was removed from several Python source files, streamlining the codebase.
• The author confirmed adherence to contributing guidelines, including running formatting tools and updating documentation.

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 significantly improves the documentation by adding comprehensive guides on checkpointing, metrics tracking, allocation modes, and agent workflows, and reorganizing existing tutorials. It also includes a cleanup by removing the unused AREAL_RECOVER_RUN environment variable. However, a security audit identified critical vulnerabilities, including multiple instances of Command Injection in the local and Slurm launchers, and an Insecure Deserialization vulnerability in the recovery handler. These issues stem from the use of unsanitized data from external sources in sensitive sinks such as subprocess.Popen(shell=True), bash script templates, and pickle.load(). Additionally, minor suggestions have been made to fix typos and improve formatting for consistency and clarity in the new documentation. It is strongly recommended to sanitize all external inputs and avoid using insecure execution patterns to mitigate these security risks.

I am having trouble creating individual review comments. Click here to see my feedback.

areal/launcher/local.py (370)

The LocalLauncher.submit_array method (lines 141-148) constructs a shell command string by concatenating environment variables and the command itself, and then executes it using subprocess.Popen with shell=True. Since some environment variables (like AREAL_LLM_SERVER_ADDRS on line 369) are sourced from an external name resolution service (line 349), an attacker who can manipulate the name resolution service can achieve arbitrary code execution on the local machine. Although this PR removes one environment variable from the dictionary, the underlying command injection vulnerability remains in the modified block.

areal/launcher/slurm.py (612)

The SlurmLauncher constructs an sbatch script by inserting environment variables and command strings directly into a bash template. This script is then written to a file and executed. Values sourced from the name resolution service (e.g., llm_addrs on line 553) are joined into AREAL_LLM_SERVER_ADDRS and inserted into the script without sanitization. An attacker manipulating the name resolution service can inject malicious bash commands into the sbatch script, leading to Remote Code Execution (RCE) on the Slurm cluster. This PR modifies the environment variable dictionary used in this process, but does not address the lack of sanitization.

areal/scheduler/slurm.py (404)

Environment variables prepared in _prepare_worker_specs are later used to construct sbatch scripts in SlurmScheduler._generate_sbatch_script. These variables are not sanitized before being inserted into the bash script template. An attacker who can control the environment or the configuration can inject arbitrary commands into the sbatch script. This PR modifies the environment variable preparation logic but leaves the injection risk unaddressed.

areal/utils/recover.py (228-229)

The RecoverInfo.load method uses pickle.load() to deserialize dataloader state from a file on disk (line 116). The path to this file is constructed from the experiment and trial names, typically located on a shared filesystem. If an attacker can write to this filesystem, they can replace the pickle file with a malicious one, leading to arbitrary code execution when the training run is recovered. By removing the check for the AREAL_RECOVER_RUN environment variable (lines 228-229), this PR expands the attack surface by allowing the vulnerable load method to be called more easily during initialization.

docs/best_practices/workflow.md (17)

This code snippet uses AsyncOpenAI, workflow_context, and aiofiles without importing them. To make the example clearer and more self-contained for users, it's good practice to include all necessary imports.

from openai import AsyncOpenAI
from areal.infra import workflow_context
import aiofiles

# Or the `run` method in agent workflows

docs/reference/alloc_mode.md (79)

There appears to be a typo in "backed". It should likely be "backend".

Note that the internal backend configurations do not affect how AReaL allocate GPUs.

docs/reference/checkpointing.md (192-194)

The numbered list items for the recovery process are all marked with 1.. While many Markdown renderers will automatically create a sequential list (1, 2, 3), it's better practice to number them sequentially in the source for clarity and to ensure consistent rendering across all platforms.

1. `RecoverHandler.load()` restores all saved state (if any)
2. Training continues from `last_step_info.next().global_step`
3. Inference engine weights are synchronized to match recovered state

docs/reference/metrics_tracking.md (34)

Using bold text for subheadings (**Characteristics:**) is less conventional than using markdown heading syntax (e.g., ### Characteristics). Using proper headings improves the document's structure, makes it easier to navigate, and allows for direct linking to sections. This applies to other similar subheadings in this file (lines 41, 65, 94, 101, 129).

### Characteristics

docs/reference/rollout_workflow.md (154)

The code example for implementing a custom workflow uses uuid and ModelRequest without importing them. This could lead to confusion or copy-paste errors for users. It's a good practice to include all necessary imports in documentation examples to make them self-contained and runnable.

from areal.api.workflow_api import RolloutWorkflow
from areal.api.io_struct import ModelRequest
import uuid

docs/tutorial/gsm8k_grpo.md (274)

There's a typo in the word "perferrable". It should be "preferable".

`agenerate` API. It is preferable if you want more fine-grained control over token IDs.

[rchardx]

LGTM