From 0f296c22ca65c149b49cdbf7a750082a6c0b55c7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=8D=9A=E6=83=9F?= Date: Mon, 2 Feb 2026 20:57:46 +0800 Subject: [PATCH] chore(infra): consolidate infrastructure and add create-pr command MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Consolidate scattered infrastructure modules into a unified `areal/infra` package for better organization and discoverability. Key changes: - Rename `areal/core` → `areal/infra` to better reflect infrastructure role - Move `areal/controller` → `areal/infra/controller` for logical grouping - Move `areal/platforms` → `areal/infra/platforms` to colocate platform abstractions - Move vLLM extensions from `areal/thirdparty/vllm` → `areal/engine/vllm_ext` to align with other engine extensions - Add `.claude/commands/create-pr.md` for automated PR creation workflow - Add launcher-scheduler expert agent documentation - Update all imports across 135 files to reflect new structure --- .claude/agents/launcher-scheduler-expert.md | 4 +- .claude/commands/create-pr.md | 692 ++++++++++++++++++ .claude/plans/wobbly-kindling-quill.md | 202 +++++ AGENTS.md | 20 +- CLAUDE.md | 16 +- CONTRIBUTING.md | 2 +- areal/__init__.py | 18 + areal/api/cli_args.py | 4 +- areal/api/engine_api.py | 2 +- areal/api/io_struct.py | 2 +- areal/engine/core/train_engine.py | 2 +- areal/engine/fsdp_engine.py | 4 +- areal/engine/megatron_engine.py | 4 +- areal/engine/ppo/actor.py | 2 +- areal/engine/ppo/critic.py | 2 +- areal/engine/rw/rw_engine.py | 4 +- areal/engine/sft/lm_engine.py | 2 +- areal/engine/sglang_remote.py | 6 +- .../vllm_ext}/areal_vllm_server.py | 0 .../vllm_ext}/vllm_worker_extension.py | 2 +- areal/engine/vllm_remote.py | 6 +- areal/experimental/engine/archon_engine.py | 4 +- areal/experimental/openai/proxy/workflow.py | 2 +- areal/experimental/trainer/rl.py | 4 +- areal/experimental/trainer/sft.py | 2 +- areal/experimental/workflow/multi_turn_v2.py | 2 +- areal/{core => infra}/__init__.py | 9 + areal/{core => infra}/async_task_runner.py | 0 areal/{ => infra}/controller/__init__.py | 4 +- .../controller/rollout_callback.py | 0 .../controller/rollout_controller.py | 5 +- .../controller/train_controller.py | 5 +- areal/{core => infra}/dist_rollout.py | 2 +- areal/{ => infra}/platforms/__init__.py | 0 areal/{ => infra}/platforms/cpu.py | 0 areal/{ => infra}/platforms/cuda.py | 0 areal/{ => infra}/platforms/npu.py | 0 areal/{ => infra}/platforms/platform.py | 0 areal/{ => infra}/platforms/unknown.py | 0 areal/{core => infra}/remote_inf_engine.py | 4 +- areal/{core => infra}/staleness_manager.py | 0 areal/{core => infra}/workflow_context.py | 0 areal/{core => infra}/workflow_executor.py | 8 +- areal/launcher/local.py | 2 +- areal/launcher/ray.py | 2 +- areal/launcher/sglang_server.py | 2 +- areal/launcher/slurm.py | 2 +- areal/launcher/vllm_server.py | 2 +- areal/models/mcore/hf_load.py | 2 +- areal/models/mcore/hf_save.py | 2 +- areal/scheduler/local.py | 2 +- areal/scheduler/ray.py | 2 +- areal/scheduler/rpc/ray_rpc_server.py | 2 +- areal/scheduler/rpc/rpc_server.py | 2 +- .../archon/test_checkpoint_e2e.py | 2 +- .../archon/test_distributed_cp.py | 2 +- .../archon/test_distributed_dp.py | 2 +- .../archon/test_distributed_ep.py | 2 +- .../archon/test_distributed_etp.py | 2 +- .../archon/test_distributed_pp.py | 2 +- .../archon/test_distributed_tp.py | 2 +- .../tests/experimental/archon/test_forward.py | 2 +- areal/tests/experimental/archon/test_grpo.py | 2 +- .../experimental/archon/test_parallel_dims.py | 2 +- .../archon/test_qwen3_parallelize.py | 2 +- .../archon/test_varlen_attention.py | 2 +- .../archon/torchrun/run_cp_forward.py | 2 +- .../archon/torchrun/run_forward.py | 2 +- .../archon/torchrun/run_tp_forward.py | 2 +- .../archon/torchrun/run_vs_fsdp.py | 2 +- areal/tests/experimental/archon/utils.py | 2 +- .../openai/test_proxy_integration.py | 2 +- areal/tests/fp8/test_fp8_bf16_comparison.py | 2 +- areal/tests/fp8/test_fp8_conversion.py | 2 +- areal/tests/test_async_task_runner.py | 2 +- areal/tests/test_data_redistribution.py | 2 +- areal/tests/test_examples.py | 2 +- areal/tests/test_fsdp_dcp.py | 2 +- .../tests/test_fsdp_memory_efficient_lora.py | 2 +- areal/tests/test_fsdp_ulysses_forward.py | 2 +- areal/tests/test_fsdp_ulysses_train_batch.py | 2 +- areal/tests/test_megatron_engine.py | 2 +- .../tests/test_megatron_engine_distributed.py | 2 +- areal/tests/test_offload.py | 2 +- .../test_packed_vs_padded_consistency.py | 2 +- areal/tests/test_perf_tracer.py | 2 +- areal/tests/test_rollout_controller.py | 2 +- areal/tests/test_staleness_manager.py | 2 +- areal/tests/test_train_controller.py | 2 +- areal/tests/test_train_engine.py | 2 +- areal/tests/test_tree_training.py | 2 +- areal/tests/test_ulysses.py | 2 +- areal/tests/test_ulysses_all_to_all.py | 2 +- areal/tests/test_vocab_parallel.py | 2 +- areal/tests/torchrun/redistribute.py | 4 +- .../torchrun/run_fsdp_dcp_distributed.py | 2 +- .../run_fsdp_memory_efficient_lora.py | 2 +- .../torchrun/run_fsdp_ulysses_forward.py | 2 +- .../torchrun/run_fsdp_ulysses_train_batch.py | 2 +- areal/tests/torchrun/run_lock.py | 2 +- .../run_megatron_engine_distributed.py | 2 +- areal/tests/torchrun/run_ulysses.py | 2 +- .../tests/torchrun/run_ulysses_all_to_all.py | 2 +- areal/tests/torchrun/run_vocab_parallel.py | 2 +- areal/tools/profile_archon.py | 2 +- areal/tools/profile_engines.py | 2 +- areal/tools/profile_fsdp.py | 2 +- areal/utils/data.py | 2 +- areal/utils/fp8/tensor_helper.py | 2 +- areal/utils/fsdp/__init__.py | 2 +- areal/utils/fsdp/grad.py | 2 +- areal/utils/fsdp/multi_tensor_apply.py | 2 +- areal/utils/fsdp/parallel.py | 2 +- areal/utils/functional/vocab_parallel.py | 2 +- areal/utils/megatron_checkpointer.py | 2 +- areal/utils/recover.py | 2 +- areal/utils/saver.py | 2 +- areal/utils/seeding.py | 2 +- areal/utils/stats_tracker.py | 2 +- areal/utils/timeutil.py | 2 +- areal/workflow/multi_turn.py | 2 +- areal/workflow/rlvr.py | 2 +- areal/workflow/vision_rlvr.py | 2 +- docs/best_practices/perf_profiling.md | 6 +- docs/cli_reference.md | 2 +- examples/alignment/hhrlhf_rw.py | 2 +- examples/camel/train.py | 2 +- examples/countdown/train.py | 2 +- examples/lora/gsm8k_grpo_lora.py | 4 +- examples/lora/gsm8k_grpo_lora_vllm.py | 2 +- examples/multi_turn_math/gsm8k_rl_mt.py | 2 +- examples/openai_agents/train_agents.py | 2 +- .../search_agent/tongyi_deepresearch/train.py | 2 +- examples/tir/tir_workflow.py | 2 +- recipe/AEnt/gsm8k_aent_grpo.py | 2 +- 135 files changed, 1087 insertions(+), 152 deletions(-) create mode 100644 .claude/commands/create-pr.md create mode 100644 .claude/plans/wobbly-kindling-quill.md rename areal/{thirdparty/vllm => engine/vllm_ext}/areal_vllm_server.py (100%) rename areal/{thirdparty/vllm => engine/vllm_ext}/vllm_worker_extension.py (99%) rename areal/{core => infra}/__init__.py (57%) rename areal/{core => infra}/async_task_runner.py (100%) rename areal/{ => infra}/controller/__init__.py (52%) rename areal/{ => infra}/controller/rollout_callback.py (100%) rename areal/{ => infra}/controller/rollout_controller.py (99%) rename areal/{ => infra}/controller/train_controller.py (99%) rename areal/{core => infra}/dist_rollout.py (99%) rename areal/{ => infra}/platforms/__init__.py (100%) rename areal/{ => infra}/platforms/cpu.py (100%) rename areal/{ => infra}/platforms/cuda.py (100%) rename areal/{ => infra}/platforms/npu.py (100%) rename areal/{ => infra}/platforms/platform.py (100%) rename areal/{ => infra}/platforms/unknown.py (100%) rename areal/{core => infra}/remote_inf_engine.py (99%) rename areal/{core => infra}/staleness_manager.py (100%) rename areal/{core => infra}/workflow_context.py (100%) rename areal/{core => infra}/workflow_executor.py (99%) diff --git a/.claude/agents/launcher-scheduler-expert.md b/.claude/agents/launcher-scheduler-expert.md index a6ca614287..c6dd1d341f 100644 --- a/.claude/agents/launcher-scheduler-expert.md +++ b/.claude/agents/launcher-scheduler-expert.md @@ -95,8 +95,8 @@ Critical utilities in `areal/utils/launcher.py`: ## Best Practices & Common Pitfalls - Use `areal.utils.logging.getLogger("LauncherName")` for logging -> not `print()` -- Query `areal.platforms.current_platform` for device information -> not hard-coded GPU - indices or direct `torch.cuda` calls +- Query `areal.infra.platforms.current_platform` for device information -> not + hard-coded GPU indices or direct `torch.cuda` calls - Use `areal.utils.name_resolve` for multi-node service discovery -> not direct IP/hostname assumptions - Raise specific exceptions from `areal.scheduler.exceptions` -> not generic exception diff --git a/.claude/commands/create-pr.md b/.claude/commands/create-pr.md new file mode 100644 index 0000000000..7974fcd295 --- /dev/null +++ b/.claude/commands/create-pr.md @@ -0,0 +1,692 @@ +--- +name: create-pr +description: Rebase from the latest `origin/main`, squash the commits from it, and then create a PR on github with intelligent commit messages based on staged changes. Invoke with /create-pr. +--- + +# Create Pull Request + +Rebase from the latest `origin/main`, squash commits, and create a PR on GitHub with an +intelligent title and description. + +## Usage + +``` +/create-pr [--draft] [--base ] +``` + +**Arguments:** + +- `--draft`: Create as draft PR +- `--base `: Target branch (default: `main`) + +## Workflow + +### Step 1: Verify Prerequisites + +```bash +# Check current branch +git branch --show-current + +# Check if on main/master (should NOT be) +if [[ $(git branch --show-current) == "main" || $(git branch --show-current) == "master" ]]; then + echo "ERROR: Cannot create PR from main/master branch" + exit 1 +fi + +# Check for uncommitted changes +git status --short + +# Ensure gh CLI is available +gh --version +``` + +**Action:** If there are uncommitted changes, stop, and then ask user to commit or stash +them first. + +### Step 2: Check for Existing PR + +```bash +# Check if PR already exists for current branch +gh pr view --json number,title,url 2>/dev/null || echo "No existing PR" +``` + +**Handle Existing PR:** + +- If PR exists, inform user and ask permission to force-update it +- Warn that this will rewrite the commit history and PR description +- If user declines, abort the process + +### Step 3: Fetch and Rebase + +```bash +# Fetch latest from origin +git fetch origin main + +# Check divergence +git log --oneline HEAD ^origin/main + +# Non-interactive rebase onto origin/main +git rebase origin/main +``` + +**Handle Conflicts:** If rebase fails due to conflicts, abort and let user handle rebase +manually: + +```bash +# On rebase failure, abort automatically +git rebase --abort + +# Inform user to resolve conflicts manually +echo "Rebase failed due to conflicts. Please resolve manually and retry /create-pr" +exit 1 +``` + +### Step 4: Squash Commits into Single Commit + +After successful rebase, squash all commits since `origin/main` into a single commit: + +```bash +# Count commits to squash +git rev-list --count origin/main..HEAD + +# Soft reset to origin/main (keeps changes staged) +git reset --soft origin/main + +# Generate commit message using /gen-commit-msg logic +# See .claude/commands/gen-commit-msg.md for message generation rules +``` + +**Generate Commit Message** (following `/gen-commit-msg` format): + +1. Analyze staged changes: + + ```bash + git diff --cached --name-only + git diff --cached + ``` + +1. Categorize changes (feat/fix/docs/refactor/test/chore/perf) + +1. Determine scope from changed files (workflow/engine/reward/dataset/api/docs/etc.) + +1. Generate message in format: + + ``` + (): + + + + [Optional sections:] + Key changes: + - change 1 + - change 2 + + Refs: #123, #456 + ``` + +1. Commit with generated message: + + ```bash + git commit -m "$(cat <<'EOF' + + EOF + )" + ``` + +### Step 5: Analyze Combined Changes + +After squashing into a single commit: + +```bash +# Get all changes since origin/main +git diff origin/main...HEAD --name-only + +# Get full diff content +git diff origin/main...HEAD + +# Check commit history +git log --oneline origin/main..HEAD +``` + +**Categorize Changes:** + +Follow same categorization as `/gen-commit-msg`: + +| Type | When to Use | +| ---------- | ------------------------------- | +| `feat` | New feature or capability | +| `fix` | Bug fix | +| `docs` | Documentation only | +| `refactor` | Code change without feature/fix | +| `test` | Adding or fixing tests | +| `chore` | Build, deps, config changes | +| `perf` | Performance improvement | + +**Determine Scope:** + +Infer from changed files: + +- `areal/workflow/` → `workflow` +- `areal/engine/` → `engine` +- `areal/reward/` → `reward` +- `areal/dataset/` → `dataset` +- `areal/api/` → `api` +- `areal/utils/` → `utils` +- `areal/infra/` → `infra` +- `docs/` → `docs` +- `examples/` → `examples` +- Multiple areas → omit scope or use broader term + +### Step 6: Generate PR Title and Description + +**PR Title Format:** + +``` +(): +``` + +**Rules:** + +- Keep under 70 characters +- Use imperative mood +- No period at end +- Mirror commit message style + +**PR Description Format:** + +MUST strictly follow the [GitHub PR template](../../.github/PULL_REQUEST_TEMPLATE.md): + +```markdown +## Description + + + +## Related Issue + + +Fixes #(issue) + +## 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](../CONTRIBUTING.md) +- [ ] 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):** + + + +## Additional Context + + +``` + +**How to Fill the Template:** + +1. **Description**: 2-4 sentences explaining what this PR does and why +1. **Related Issue**: Link to issue (search for related issues if exists) +1. **Type of Change**: Mark ONE primary type with `[x]` +1. **Checklist**: Mark completed items with `[x]`, leave uncompleted as `[ ]` +1. **Breaking Change Details**: Only if breaking changes checkbox is marked +1. **Additional Context**: Any extra info, related PRs, performance numbers, etc. + +### Step 7: Push and Create/Update PR + +Show preview to user: + +``` +───────────────────────────────────────────────── +Branch: feat/vision-rlvr → main + +PR Title: +feat(workflow): add vision support to RLVR + +PR Description: +## Description + +Add VisionRLVRWorkflow for vision-language RL training. Supports image inputs +alongside text prompts and integrates with existing RLVR pipeline. + +## Related Issue + +Fixes #789 + +## Type of Change + +- [ ] Bug fix (non-breaking change that fixes an issue) +- [x] 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 + +- [x] I have read the [Contributing Guide](../CONTRIBUTING.md) +- [x] I have run formatting tools (pre-commit or manual) +- [ ] I have run relevant unit tests and they pass +- [x] I have added tests for new functionality +- [x] I have updated documentation if needed +- [x] My branch is up to date with main +- [ ] This PR introduces breaking changes (if yes, fill out details below) +- [x] 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 + +Requires Pillow>=10.0.0 for image processing. + +Files changed: +- `areal/workflow/vision_rlvr.py`: New VisionRLVRWorkflow class +- `areal/api/workflow_api.py:45`: Add vision config fields +- `examples/vision_rlvr.py`: Example training script +- `docs/workflows/vision.md`: Documentation + +───────────────────────────────────────────────── + +Commands to execute: +1. git push -u origin feat/vision-rlvr +2. gh pr create --title "..." --body "..." [--draft] +───────────────────────────────────────────────── +``` + +**Confirm with user**, then execute: + +```bash +# Force push branch to remote (required after squash) +git push -f -u origin $(git branch --show-current) + +# Create or edit PR using gh CLI with GitHub template format +# If PR exists, use 'gh pr edit' instead of 'gh pr create' +if gh pr view &>/dev/null; then + # Update existing PR + gh pr edit \ + --title "feat(workflow): add vision support to RLVR" \ + --body "$(cat <<'EOF' +[PR description here] +EOF +)" +else + # Create new PR + gh pr create \ + --base main \ + --title "feat(workflow): add vision support to RLVR" \ + --body "$(cat <<'EOF' +## Description + +Add VisionRLVRWorkflow for vision-language RL training. Supports image inputs +alongside text prompts and integrates with existing RLVR pipeline. + +## Related Issue + +Fixes #789 + +## Type of Change + +- [ ] Bug fix (non-breaking change that fixes an issue) +- [x] 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 + +- [x] I have read the [Contributing Guide](../CONTRIBUTING.md) +- [x] I have run formatting tools (pre-commit or manual) +- [ ] I have run relevant unit tests and they pass +- [x] I have added tests for new functionality +- [x] I have updated documentation if needed +- [x] My branch is up to date with main +- [ ] This PR introduces breaking changes (if yes, fill out details below) +- [x] 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 + +Requires Pillow>=10.0.0 for image processing. + +Files changed: +- `areal/workflow/vision_rlvr.py`: New VisionRLVRWorkflow class +- `areal/api/workflow_api.py:45`: Add vision config fields +- `examples/vision_rlvr.py`: Example training script +- `docs/workflows/vision.md`: Documentation +EOF +)" +fi +``` + +Add `--draft` flag if requested. + +**Capture PR URL** and display to user: + +``` +✓ PR created/updated successfully! +https://github.com/inclusionAI/AReaL/pull/123 +``` + +## Examples + +### Example 1: Feature PR + +**Changes:** New dataset loader for MATH dataset + +**PR Title:** + +``` +feat(dataset): add MATH dataset loader +``` + +**PR Description:** + +```markdown +## Description + +Add MATHDataset loader for mathematics problem solving with LaTeX rendering and +symbolic math parsing. Includes reward function for automatic answer verification +and full test coverage. + +## Related Issue + +Fixes #456 + +## Type of Change + +- [ ] Bug fix (non-breaking change that fixes an issue) +- [x] 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 + +- [x] I have read the [Contributing Guide](../CONTRIBUTING.md) +- [x] I have run formatting tools (pre-commit or manual) +- [x] I have run relevant unit tests and they pass +- [x] I have added tests for new functionality +- [x] I have updated documentation if needed +- [x] My branch is up to date with main +- [ ] This PR introduces breaking changes (if yes, fill out details below) +- [x] 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 + +Dataset requires ~500MB download on first use. Added comprehensive test suite +covering all 12,500 problems with >95% reward function accuracy. + +Files changed: +- `areal/dataset/math.py`: New MATHDataset class +- `areal/reward/math_reward.py`: Symbolic math reward function +- `examples/math_training.py`: Training script +- `docs/datasets/math.md`: Dataset documentation +- `tests/test_math_dataset.py`: Unit tests +``` + +### Example 2: Bug Fix PR + +**Changes:** Fix memory leak in ArchonEngine + +**PR Title:** + +``` +fix(engine): resolve memory leak in ArchonEngine rollout +``` + +**PR Description:** + +```markdown +## Description + +Fix memory leak during ArchonEngine rollout phase by clearing cached activations +after each batch and moving tensors to CPU before deletion. Reduces memory usage +by ~2GB per rollout iteration. + +## Related Issue + +Fixes #872 + +## Type of Change + +- [x] 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 + +- [x] I have read the [Contributing Guide](../CONTRIBUTING.md) +- [x] I have run formatting tools (pre-commit or manual) +- [x] I have run relevant unit tests and they pass +- [x] I have added tests for new functionality +- [ ] I have updated documentation if needed +- [x] 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 + +Tested with 100 rollout iterations without OOM. Memory usage stable at 8GB +(previously would grow to 10GB+). Output correctness validated unchanged. + +Backported to v0.5.x branch. + +Files changed: +- `areal/engine/archon.py:234`: Add explicit cache clearing +- `areal/engine/archon.py:456`: Move tensor to CPU before deletion +- `tests/test_archon_memory.py`: Add memory leak regression test +``` + +### Example 3: Breaking Change PR + +**Changes:** Refactor reward API for better extensibility + +**PR Title:** + +``` +refactor(reward): simplify reward function interface +``` + +**PR Description:** + +```markdown +## Description + +Simplify reward function API from 4 methods to 2 by consolidating compute and +compute_batch into a single batched interface. Improves type hints and +documentation. All existing reward functions updated. + +## Related Issue + +Fixes #901 + +## Type of Change + +- [ ] Bug fix (non-breaking change that fixes an issue) +- [ ] New feature (non-breaking change that adds functionality) +- [x] 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 + +- [x] I have read the [Contributing Guide](../CONTRIBUTING.md) +- [x] I have run formatting tools (pre-commit or manual) +- [x] I have run relevant unit tests and they pass +- [ ] I have added tests for new functionality +- [x] I have updated documentation if needed +- [x] My branch is up to date with main +- [x] This PR introduces breaking changes (if yes, fill out details below) +- [x] 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):** + +Old `compute_batch` method is deprecated and will be removed in v0.7.0. + +See migration guide in `docs/migration/reward_api.md` for details. + +## Additional Context + +All existing tests pass. Performance unchanged at 10k rewards/sec. Backward +compatibility warnings added for deprecated methods. + +Files changed: +- `areal/api/reward_api.py:12`: Consolidate compute/compute_batch +- `areal/reward/gsm8k.py`: Update to new API +- `areal/reward/code_reward.py`: Update to new API +- `areal/reward/geometry3k.py`: Update to new API +- `docs/customization/reward.md`: Update documentation +- `docs/migration/reward_api.md`: Migration guide +- `examples/custom_reward.py`: Update example +``` + +## Error Handling + +### Rebase Conflicts + +If rebase fails: + +1. Show conflict files +1. Provide resolution instructions +1. Wait for user to resolve +1. After resolution, continue with squashing step +1. Offer to abort rebase if needed: `git rebase --abort` + +### Squash Failures + +If squash/commit fails: + +1. Check if there are changes to commit: `git status` +1. Verify no conflicts remain: `git diff --cached` +1. If needed, abort and return to pre-rebase state + +### Push Failures + +If force push fails: + +1. Verify remote branch exists +1. Check GitHub authentication: `gh auth status` +1. Confirm branch protection rules allow force push +1. Provide manual push instructions if needed + +### PR Creation/Update Failures + +If `gh pr create` or `gh pr edit` fails: + +1. Check if PR already exists: `gh pr view` +1. Verify GitHub authentication: `gh auth status` +1. Check for branch protection rules +1. Provide manual PR creation/update link + +## Safety Checks + +**Before Starting:** + +- Confirm no uncommitted changes +- Confirm not on main/master branch +- Check for existing PR and get user permission to overwrite if exists +- Backup branch: `git branch backup/$(git branch --show-current)-$(date +%s)` + +**Before Rebase:** + +- Fetch latest from origin +- Show divergence summary + +**Before Squash:** + +- Show commits that will be squashed +- Confirm user wants to proceed + +**Before Force Push:** + +- **CRITICAL**: Warn user that force push will rewrite history +- Show current commit that will replace remote history +- Confirm branch name +- If PR exists, emphasize that PR history will be rewritten + +**Before PR Creation/Update:** + +- Show full preview of title/description +- Confirm target branch +- If updating existing PR, show what will change + +______________________________________________________________________ + + diff --git a/.claude/plans/wobbly-kindling-quill.md b/.claude/plans/wobbly-kindling-quill.md new file mode 100644 index 0000000000..a0625ba3e5 --- /dev/null +++ b/.claude/plans/wobbly-kindling-quill.md @@ -0,0 +1,202 @@ +# Plan: Rewrite Rollout Section with Accurate Async Details + +## Background + +The user wants to: + +1. Remove lines 236-248 (redundant TrainController/RolloutController subsections under + PPOTrainer) +1. Rewrite the "Rollout: Generating Training Data" section to accurately explain + asynchronous generation + +**Key insight from user**: "Asynchronous" here means: + +- Python async/await semantics +- **Background thread** (`BatchTaskDispatcher`) submitting rollout requests concurrently +- **Separate GPU process** (SGLang/vLLM server) handling requests on GPU +- **Overlap generation with training**: Rollout and training happen simultaneously + +## Current Issues + +**Lines 236-248**: Redundant subsections after we already have H2 sections for Train and +Rollout + +- "TrainController: RPC Mechanism" +- "RolloutController: Async Workflow Execution" + +**Lines 287-406**: Current async explanation is too simplistic + +- Doesn't explain the background thread architecture +- Doesn't explain how RPC server receives and queues requests +- Doesn't explain GPU process separation (SGLang/vLLM runs independently) +- Missing the overlap with training explanation + +## Key Code Insights + +From `areal/infra/controller/rollout_controller.py`: + +1. **RolloutController** creates workers via scheduler +1. Each worker runs an RPC server (`areal/scheduler/rpc/rpc_server.py`) +1. RPC server launches SGLang/vLLM as separate subprocess +1. **BatchTaskDispatcher** (`_dispatcher`) runs in background thread +1. Dispatcher submits tasks via `_create_submit_callback()` which: + - Chooses worker (round-robin) + - Calls `scheduler.async_call_engine(worker_id, "submit", ...)` + - Returns immediately (non-blocking) + - Worker's `RemoteInfEngine.submit()` queues work to SGLang/vLLM + - Callback received when trajectory completes + +From `areal/infra/remote_inf_engine.py`: + +1. **RemoteInfEngine** wraps SGLang/vLLM HTTP servers +1. `submit()` method: + - Resolves workflow from string + - Calls `workflow_executor.submit()` which queues the task + - Returns task_id immediately (non-blocking) +1. **WorkflowExecutor** runs in background thread +1. When trajectory completes, sends HTTP callback to RolloutController + +From `areal/scheduler/rpc/rpc_server.py`: + +1. **Flask server** with `/create_engine`, `/call` endpoints +1. **Engine thread** (`_engine_thread`) runs all engine operations serially +1. SGLang/vLLM launched as subprocess via `backend.launch_server()` +1. Separate GPU process handles actual generation + +## Architecture Diagram Needed + +``` +Controller Process Worker Process (RPC Server) GPU Process +────────────────── ─────────────────────────── ─────────── +RolloutController Flask HTTP Server (CPU) SGLang/vLLM + │ │ │ + └─> BatchTaskDispatcher /call endpoint Inference Engine + (background thread) │ │ + │ └─> Engine Thread │ + ├─ submit task 1 └─> RemoteInfEngine │ + │ (HTTP POST) └─> workflow_executor │ + │ └─> submit() ───────────>│ + ├─ submit task 2 Generate + │ (HTTP POST) tokens + │ │ + ├─ submit task 3 HTTP Callback <────────┘ + │ (trajectory) + │ ┌──────────────┘ + └─ collect results <────────┘ + +Meanwhile... +TrainController Training Worker (RPC Server) + │ │ + └─> compute_logp(batch) ────────────>│ Forward pass + └─> ppo_update(batch) ──────────────>│ Backward pass + +Key: Generation and training happen SIMULTANEOUSLY on different GPUs +``` + +## Implementation Plan + +### 1. Remove Lines 236-248 + +Delete the redundant subsections: + +- "### TrainController: RPC Mechanism" (L236-246) +- "### RolloutController: Async Workflow Execution" (L248) + +These are redundant since we have separate H2 sections for training and rollout. + +### 2. Rewrite "Rollout: Generating Training Data" (L250-406) + +**New structure:** + +**L250: Keep H2 header** + +```markdown +## Rollout: Generating Training Data +``` + +**L252-266: Keep Workflow Specification (minor edits)** + +- Current content is good, just ensure it's concise + +**L268-285: Keep RLVRWorkflow overview** + +- Current 4-step description is good + +**L287+: REWRITE "Asynchronous Rollout Collection"** + +New content should explain: + +1. **Architecture Overview** (with diagram): + + - Controller process: RolloutController + BatchTaskDispatcher (background thread) + - Worker process: RPC server (Flask HTTP on CPU) + Engine thread + - GPU process: SGLang/vLLM server (separate subprocess) + - HTTP callbacks for result collection + +1. **Three Levels of Concurrency**: + + - **Level 1 (Controller)**: BatchTaskDispatcher background thread submits requests + - **Level 2 (Worker RPC)**: Flask server accepts concurrent HTTP requests, Engine + thread processes serially + - **Level 3 (GPU)**: SGLang/vLLM subprocess handles multiple concurrent generation + requests + +1. **Request Flow**: + + ``` + 1. Controller: actor.prepare_batch() calls rollout.prepare_batch() + 2. RolloutController.prepare_batch(): + - Creates task_input_generator from dataloader + - Dispatcher submits tasks to workers (round-robin) + 3. Worker RPC receives HTTP POST /call (method="submit"): + - Deserializes workflow string + kwargs + - Engine thread runs RemoteInfEngine.submit() + - Queues work to SGLang/vLLM subprocess + - Returns task_id (non-blocking) + 4. SGLang/vLLM subprocess: + - Processes requests from queue + - Generates tokens on GPU + - Returns trajectory via HTTP callback + 5. Controller receives callback: + - BatchTaskDispatcher collects results + - Waits for batch_size accepted trajectories + - Concatenates and returns + ``` + +1. **Overlap with Training**: + + - While training workers compute gradients, rollout workers continue generating + - BatchTaskDispatcher maintains 2+ batches of pending requests + - Staleness control ensures generated data isn't too old + +1. **Staleness Management**: + + - Brief explanation of version tracking + - `max_head_offpolicyness` config + - Pause/resume mechanism during weight sync + +### 3. Sections After Rollout + +Keep and verify: + +- "Training: Distributed Computing with Controllers" (should focus on TrainController + details) +- "Training Methods: What Happens on Workers" +- "Weight Synchronization" +- "Monitoring and Utilities" + +## Files to Modify + +1. **`docs/lite/gsm8k_grpo.md`** + - Remove: Lines 236-248 + - Rewrite: Lines 287-406 (Asynchronous Rollout Collection section) + +## Verification + +After completion: + +1. Run `mdformat docs/lite/gsm8k_grpo.md` +1. Run `pre-commit run --files docs/lite/gsm8k_grpo.md` +1. Verify all code references are accurate +1. Check that async architecture is correctly explained +1. Ensure diagram accurately reflects the three-process architecture diff --git a/AGENTS.md b/AGENTS.md index a2cf4f2684..55327acab2 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -22,14 +22,23 @@ When unsure, leave a `TODO(agent)` comment and note the constraint in your respo ## Repository map +- `.claude/` — Configuration directory for AI coding agents: + - `.claude/agents/` — Specialized agent definitions (planner, code-verifier, + simple-code-reviewer, engine experts, algorithm-expert, launcher-scheduler-expert) + - `.claude/commands/` — User-invocable commands (create-pr, gen-commit-msg, pr-review) + - `.claude/rules/` — Project-wide code quality standards (api-config, code-style, + distributed, testing) + - `.claude/skills/` — Guided development workflows (add-dataset, add-workflow, + add-reward, add-unit-tests, debug-distributed) + - `.claude/hooks/` — Event-driven automation hooks + - `.claude/data/` — Agent working data and caches + - `.claude/plans/` — Implementation plans generated by planner agent - `areal/` — Core Python package housing APIs, controllers, engines, workflows, and shared utilities: - `areal/api/` — Contracts for workflows, engines, schedulers, IO structs, and CLI/config dataclasses. - - `areal/controller/` — Distributed batching and controller-side dataset packing - helpers. - - `areal/core/` — Async orchestration primitives for task runners, remote inference, - and workflow execution. + - `areal/infra/` — Single controller implementation, async orchestration primitives, + and hardware/platform abstractions for CPU/GPU/NPU runtimes. - `areal/dataset/` — Stateful dataset loaders (GSM8K, Geometry3K, CLEVR, HH-RLHF, TORL, etc.) and utilities that feed rollout jobs safely. - `areal/engine/` — Training backends (FSDP2, Megatron, PPO, SFT, reward modeling) and @@ -40,15 +49,12 @@ When unsure, leave a `TODO(agent)` comment and note the constraint in your respo SGLang/vLLM inference server launchers and container guidance. - `areal/models/` — Model-specific adapters (Megatron-Core layers, Transformers wrappers, custom heads). - - `areal/platforms/` — Hardware/platform abstractions for CPU/GPU/NPU runtimes and - device adapters. - `areal/reward/` — Built-in reward functions (GSM8K, Geometry3K, CLEVR, etc.), math parsers, and helpers; wrap slow logic with `AsyncRewardWrapper`. - `areal/scheduler/` — Placement and allocation policies aligned with launcher configs. - `areal/tests/` — Focused unit/integration suites (many require GPUs or mocked distributed backends). - - `areal/thirdparty/` — Vendored integrations (e.g., vLLM/SGLang shims) kept in-tree. - `areal/tools/` — Developer utilities and maintenance scripts tied to the core package. - `areal/utils/` — Cross-cutting helpers for logging, tensor ops, stats tracking, diff --git a/CLAUDE.md b/CLAUDE.md index 0a60ec75b6..556392f1f7 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -99,8 +99,8 @@ uv run python docs/generate_cli_docs.py ## Extended Configuration -See `.claude/agents/`, `.claude/skills/`, and `.claude/commands/` for specialized -instructions. +See `.claude/agents/`, `.claude/skills/`, `.claude/commands/`, and `.claude/rules/` for +specialized instructions. ### Agents @@ -138,5 +138,15 @@ Skills provide step-by-step guides for common development tasks: Commands perform specific actions when invoked: -- `/pr-review` - Intelligent PR code review with dynamic agent allocation +- `/create-pr` - Rebase, squash commits, and create/update PR with intelligent messages - `/gen-commit-msg` - Generate commit messages from staged changes +- `/pr-review` - Intelligent PR code review with dynamic agent allocation + +### Rules (Code Quality Standards) + +Project-wide standards enforced across all code changes: + +- `api-config.md` - Configuration dataclass design patterns +- `code-style.md` - Coding conventions beyond pre-commit hooks +- `distributed.md` - Distributed training patterns and constraints +- `testing.md` - Testing strategy and coverage requirements diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 5f286fea1a..4c4af0cff9 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -244,7 +244,7 @@ failures. For example: ```python import pytest -from areal.platforms import current_platform +from areal.infra.platforms import current_platform # ordinary tests are supposed to run fast, and will run in CI def test_fast_operation(): diff --git a/areal/__init__.py b/areal/__init__.py index 4e59471784..ba6f673188 100644 --- a/areal/__init__.py +++ b/areal/__init__.py @@ -1,3 +1,21 @@ """AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language Reasoning""" from .version import __version__ # noqa + +from .infra import ( + TrainController, + RolloutController, + WorkflowExecutor, + StalenessManager, + workflow_context, + current_platform, +) + +__all__ = [ + "TrainController", + "RolloutController", + "WorkflowExecutor", + "StalenessManager", + "workflow_context", + "current_platform", +] diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py index 7f37e610ec..608a90983e 100644 --- a/areal/api/cli_args.py +++ b/areal/api/cli_args.py @@ -1097,7 +1097,7 @@ class vLLMConfig: no_enable_prefix_caching: bool = True gpu_memory_utilization: float = 0.9 worker_extension_cls: str = ( - "areal.thirdparty.vllm.vllm_worker_extension.VLLMWorkerExtension" + "areal.engine.vllm_ext.vllm_worker_extension.VLLMWorkerExtension" ) enable_sleep_mode: bool = False uvicorn_log_level: str = "warning" @@ -1143,7 +1143,7 @@ def build_args( @staticmethod def build_cmd_from_args(args: dict[str, Any]): - return get_py_cmd("areal.thirdparty.vllm.areal_vllm_server", args) + return get_py_cmd("areal.engine.vllm_ext.areal_vllm_server", args) @staticmethod def build_cmd( diff --git a/areal/api/engine_api.py b/areal/api/engine_api.py index 678531bb53..5ae8f687c9 100644 --- a/areal/api/engine_api.py +++ b/areal/api/engine_api.py @@ -23,7 +23,7 @@ if TYPE_CHECKING: from areal.api.workflow_api import WorkflowLike - from areal.core.workflow_executor import WorkflowExecutor + from areal.infra import WorkflowExecutor from areal.utils.data import MicroBatchList diff --git a/areal/api/io_struct.py b/areal/api/io_struct.py index a7fcc8bb37..26ab52216d 100644 --- a/areal/api/io_struct.py +++ b/areal/api/io_struct.py @@ -12,7 +12,7 @@ from areal.api.alloc_mode import AllocationMode from areal.api.cli_args import GenerationHyperparameters -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils import logging if TYPE_CHECKING: diff --git a/areal/engine/core/train_engine.py b/areal/engine/core/train_engine.py index c925cd6ec0..ef1bb3b917 100644 --- a/areal/engine/core/train_engine.py +++ b/areal/engine/core/train_engine.py @@ -10,7 +10,7 @@ import torch import torch.distributed as dist -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.data import ( MicroBatchList, pad_and_stack_tensors_along_first_dim, diff --git a/areal/engine/fsdp_engine.py b/areal/engine/fsdp_engine.py index 2adc5f768f..e172bd1085 100644 --- a/areal/engine/fsdp_engine.py +++ b/areal/engine/fsdp_engine.py @@ -52,12 +52,13 @@ WeightUpdateMeta, ) from areal.api.workflow_api import WorkflowLike -from areal.core.dist_rollout import DistRolloutCoordinator from areal.engine.core import ( aggregate_eval_losses, compute_total_loss_weight, reorder_and_pad_outputs, ) +from areal.infra.dist_rollout import DistRolloutCoordinator +from areal.infra.platforms import current_platform from areal.models.fsdp.ulysses import ( set_ulysses_sequence_parallel_group, ulysses_pad, @@ -76,7 +77,6 @@ patch_fsdp_for_tree_training, ) from areal.models.tree_attn.tree import TrieNode, build_packed_tree_batch -from areal.platforms import current_platform from areal.utils import ( logging, name_resolve, diff --git a/areal/engine/megatron_engine.py b/areal/engine/megatron_engine.py index acdc574bf0..5cd9afc79f 100644 --- a/areal/engine/megatron_engine.py +++ b/areal/engine/megatron_engine.py @@ -43,12 +43,13 @@ WeightUpdateMeta, ) from areal.api.workflow_api import WorkflowLike -from areal.core.dist_rollout import DistRolloutCoordinator from areal.engine.core import ( aggregate_eval_losses, compute_total_loss_weight, reorder_and_pad_outputs, ) +from areal.infra.dist_rollout import DistRolloutCoordinator +from areal.infra.platforms import current_platform from areal.models.mcore.hf_load import load_weights_from_hf_with_mbridge_fast from areal.models.mcore.hf_save import save_weights_to_hf_with_mbridge_fast from areal.models.mcore.registry import make_hf_and_mcore_config, make_mcore_model @@ -63,7 +64,6 @@ patch_bridge_for_tree_training, ) from areal.models.tree_attn.tree import build_packed_tree_batch -from areal.platforms import current_platform from areal.utils import logging, name_resolve, names, perf_tracer, stats_tracker from areal.utils.constants import ( DEFAULT_VECTORIZED_ALIGNMENT_BYTES, diff --git a/areal/engine/ppo/actor.py b/areal/engine/ppo/actor.py index e981a3c86e..dd2d579496 100644 --- a/areal/engine/ppo/actor.py +++ b/areal/engine/ppo/actor.py @@ -5,7 +5,7 @@ from areal.api.cli_args import MicroBatchSpec, PPOActorConfig from areal.api.engine_api import TrainEngine -from areal.controller.train_controller import TrainController +from areal.infra import TrainController from areal.utils import logging, stats_tracker from areal.utils.constants import ( PROX_APPROX_METHOD_LINEAR, diff --git a/areal/engine/ppo/critic.py b/areal/engine/ppo/critic.py index dd27c24c69..24ea6db252 100644 --- a/areal/engine/ppo/critic.py +++ b/areal/engine/ppo/critic.py @@ -5,7 +5,7 @@ from areal.api.cli_args import MicroBatchSpec, PPOCriticConfig from areal.api.engine_api import TrainEngine -from areal.controller.train_controller import TrainController +from areal.infra import TrainController from areal.utils import stats_tracker from areal.utils.data import split_padded_tensor_dict_into_mb_list from areal.utils.functional import ppo_critic_loss_fn diff --git a/areal/engine/rw/rw_engine.py b/areal/engine/rw/rw_engine.py index bc5a2bbba5..c1fe746b41 100644 --- a/areal/engine/rw/rw_engine.py +++ b/areal/engine/rw/rw_engine.py @@ -3,8 +3,8 @@ import torch from areal.api.engine_api import TrainEngine -from areal.controller.train_controller import TrainController -from areal.platforms import current_platform +from areal.infra import TrainController +from areal.infra.platforms import current_platform from areal.utils import logging, stats_tracker from areal.utils.perf_tracer import trace_perf diff --git a/areal/engine/sft/lm_engine.py b/areal/engine/sft/lm_engine.py index 79d4c3c4d9..93071d6401 100644 --- a/areal/engine/sft/lm_engine.py +++ b/areal/engine/sft/lm_engine.py @@ -3,7 +3,7 @@ import torch from areal.api.engine_api import TrainEngine -from areal.controller.train_controller import TrainController +from areal.infra import TrainController from areal.utils import stats_tracker from areal.utils.perf_tracer import trace_perf diff --git a/areal/engine/sglang_remote.py b/areal/engine/sglang_remote.py index 2df2686c06..f9c5151c90 100644 --- a/areal/engine/sglang_remote.py +++ b/areal/engine/sglang_remote.py @@ -22,10 +22,8 @@ ) from areal.api.scheduler_api import Scheduler from areal.api.workflow_api import WorkflowLike -from areal.controller import RolloutController -from areal.core import RemoteInfEngine -from areal.core.workflow_executor import WorkflowExecutor -from areal.platforms import current_platform +from areal.infra import RemoteInfEngine, RolloutController, WorkflowExecutor +from areal.infra.platforms import current_platform from areal.utils import perf_tracer, stats_tracker from areal.utils.launcher import TRITON_CACHE_PATH diff --git a/areal/thirdparty/vllm/areal_vllm_server.py b/areal/engine/vllm_ext/areal_vllm_server.py similarity index 100% rename from areal/thirdparty/vllm/areal_vllm_server.py rename to areal/engine/vllm_ext/areal_vllm_server.py diff --git a/areal/thirdparty/vllm/vllm_worker_extension.py b/areal/engine/vllm_ext/vllm_worker_extension.py similarity index 99% rename from areal/thirdparty/vllm/vllm_worker_extension.py rename to areal/engine/vllm_ext/vllm_worker_extension.py index 0757ecadb0..057b5c143c 100644 --- a/areal/thirdparty/vllm/vllm_worker_extension.py +++ b/areal/engine/vllm_ext/vllm_worker_extension.py @@ -8,7 +8,7 @@ from vllm.lora.request import LoRARequest from vllm.model_executor.model_loader import get_model_loader -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.constants import DIST_GROUP_DEFAULT_TIMEOUT from areal.utils.distributed import init_custom_process_group diff --git a/areal/engine/vllm_remote.py b/areal/engine/vllm_remote.py index cfb349c457..81443eaa47 100644 --- a/areal/engine/vllm_remote.py +++ b/areal/engine/vllm_remote.py @@ -22,10 +22,8 @@ ) from areal.api.scheduler_api import Scheduler from areal.api.workflow_api import WorkflowLike -from areal.controller import RolloutController -from areal.core import RemoteInfEngine -from areal.core.workflow_executor import WorkflowExecutor -from areal.platforms import current_platform +from areal.infra import RemoteInfEngine, RolloutController, WorkflowExecutor +from areal.infra.platforms import current_platform from areal.utils import perf_tracer, stats_tracker from areal.utils.launcher import TRITON_CACHE_PATH diff --git a/areal/experimental/engine/archon_engine.py b/areal/experimental/engine/archon_engine.py index 123e6b5175..961628c2df 100644 --- a/areal/experimental/engine/archon_engine.py +++ b/areal/experimental/engine/archon_engine.py @@ -37,7 +37,6 @@ ) from areal.api.scheduler_api import Scheduler from areal.api.workflow_api import WorkflowLike -from areal.core.dist_rollout import DistRolloutCoordinator from areal.engine.core.train_engine import ( aggregate_eval_losses, compute_total_loss_weight, @@ -66,7 +65,8 @@ ulysses_gather_output, ulysses_slice_inputs, ) -from areal.platforms import current_platform +from areal.infra.dist_rollout import DistRolloutCoordinator +from areal.infra.platforms import current_platform from areal.utils import logging, name_resolve, names, perf_tracer, stats_tracker from areal.utils.constants import DEFAULT_PAGE_SIZE_BYTES, DIST_GROUP_DEFAULT_TIMEOUT from areal.utils.data import ( diff --git a/areal/experimental/openai/proxy/workflow.py b/areal/experimental/openai/proxy/workflow.py index 532478d037..f92a51a5bb 100644 --- a/areal/experimental/openai/proxy/workflow.py +++ b/areal/experimental/openai/proxy/workflow.py @@ -10,7 +10,7 @@ import aiohttp from areal.api.workflow_api import AgentWorkflow, RolloutWorkflow -from areal.core import workflow_context +from areal.infra import workflow_context from areal.utils import logging, stats_tracker from areal.utils.perf_tracer import session_context, trace_session diff --git a/areal/experimental/trainer/rl.py b/areal/experimental/trainer/rl.py index a4b7b5a5eb..e71ff5e2d7 100644 --- a/areal/experimental/trainer/rl.py +++ b/areal/experimental/trainer/rl.py @@ -27,10 +27,10 @@ from areal.api.io_struct import FinetuneSpec, StepInfo, WeightUpdateMeta from areal.api.scheduler_api import Scheduler from areal.api.workflow_api import AgentWorkflow, WorkflowLike -from areal.controller import RolloutController from areal.engine.sglang_remote import RemoteSGLangEngine from areal.engine.vllm_remote import RemotevLLMEngine -from areal.platforms import current_platform +from areal.infra import RolloutController +from areal.infra.platforms import current_platform from areal.scheduler import LocalScheduler, RayScheduler, SlurmScheduler from areal.utils import logging, perf_tracer, seeding, stats_tracker from areal.utils.dataloader import create_dataloader diff --git a/areal/experimental/trainer/sft.py b/areal/experimental/trainer/sft.py index e04fb5a62f..c1c2b9498c 100644 --- a/areal/experimental/trainer/sft.py +++ b/areal/experimental/trainer/sft.py @@ -16,7 +16,7 @@ ) from areal.api.io_struct import FinetuneSpec, StepInfo from areal.api.scheduler_api import Scheduler -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.scheduler import LocalScheduler, RayScheduler, SlurmScheduler from areal.utils import logging, perf_tracer, seeding, stats_tracker from areal.utils.data import ( diff --git a/areal/experimental/workflow/multi_turn_v2.py b/areal/experimental/workflow/multi_turn_v2.py index f3cfa647a1..9430b496f0 100644 --- a/areal/experimental/workflow/multi_turn_v2.py +++ b/areal/experimental/workflow/multi_turn_v2.py @@ -3,11 +3,11 @@ from transformers import PreTrainedTokenizerFast +from areal import workflow_context from areal.api.cli_args import GenerationHyperparameters from areal.api.engine_api import InferenceEngine from areal.api.reward_api import AsyncRewardWrapper from areal.api.workflow_api import RolloutWorkflow -from areal.core import workflow_context from areal.experimental.openai import ArealOpenAI from areal.utils import logging, stats_tracker diff --git a/areal/core/__init__.py b/areal/infra/__init__.py similarity index 57% rename from areal/core/__init__.py rename to areal/infra/__init__.py index 28fa01204a..bbe624b13b 100644 --- a/areal/core/__init__.py +++ b/areal/infra/__init__.py @@ -1,5 +1,8 @@ """Core components for AREAL.""" +from . import workflow_context +from .controller import RolloutController, TrainController +from .platforms import Platform, current_platform, is_npu_available from .remote_inf_engine import ( RemoteInfBackendProtocol, RemoteInfEngine, @@ -16,4 +19,10 @@ "StalenessManager", "WorkflowExecutor", "check_trajectory_format", + "RolloutController", + "TrainController", + "workflow_context", + "Platform", + "current_platform", + "is_npu_available", ] diff --git a/areal/core/async_task_runner.py b/areal/infra/async_task_runner.py similarity index 100% rename from areal/core/async_task_runner.py rename to areal/infra/async_task_runner.py diff --git a/areal/controller/__init__.py b/areal/infra/controller/__init__.py similarity index 52% rename from areal/controller/__init__.py rename to areal/infra/controller/__init__.py index 45575be055..feb8c72e0d 100644 --- a/areal/controller/__init__.py +++ b/areal/infra/controller/__init__.py @@ -1,7 +1,7 @@ """Controller components for managing distributed training and inference.""" -from areal.controller.rollout_controller import RolloutController -from areal.controller.train_controller import TrainController +from .rollout_controller import RolloutController +from .train_controller import TrainController __all__ = [ "RolloutController", diff --git a/areal/controller/rollout_callback.py b/areal/infra/controller/rollout_callback.py similarity index 100% rename from areal/controller/rollout_callback.py rename to areal/infra/controller/rollout_callback.py diff --git a/areal/controller/rollout_controller.py b/areal/infra/controller/rollout_controller.py similarity index 99% rename from areal/controller/rollout_controller.py rename to areal/infra/controller/rollout_controller.py index d4ad3c7af4..3e1fb675ba 100644 --- a/areal/controller/rollout_controller.py +++ b/areal/infra/controller/rollout_controller.py @@ -29,8 +29,6 @@ ) from areal.api.scheduler_api import Job, Scheduler, Worker from areal.api.workflow_api import AgentWorkflow, RolloutWorkflow, WorkflowLike -from areal.core.staleness_manager import StalenessManager -from areal.core.workflow_executor import BatchTaskDispatcher, TaskIdGenerator from areal.scheduler.rpc.serialization import deserialize_value from areal.utils import logging, perf_tracer from areal.utils.concurrent import run_async_task @@ -39,6 +37,9 @@ from areal.utils.network import find_free_ports, gethostip from areal.utils.perf_tracer import trace_perf +from ..staleness_manager import StalenessManager +from ..workflow_executor import BatchTaskDispatcher, TaskIdGenerator + logger = logging.getLogger("RolloutController") diff --git a/areal/controller/train_controller.py b/areal/infra/controller/train_controller.py similarity index 99% rename from areal/controller/train_controller.py rename to areal/infra/controller/train_controller.py index e3aacdff86..74f387c140 100644 --- a/areal/controller/train_controller.py +++ b/areal/infra/controller/train_controller.py @@ -15,13 +15,14 @@ ) from areal.api.scheduler_api import Job, Scheduler, Worker from areal.api.workflow_api import WorkflowLike -from areal.controller.rollout_callback import RolloutCallback -from areal.controller.rollout_controller import RolloutController from areal.scheduler.rpc.rtensor import RTensor from areal.utils import logging, stats_tracker from areal.utils.concurrent import run_async_task from areal.utils.network import find_free_ports +from .rollout_callback import RolloutCallback +from .rollout_controller import RolloutController + logger = logging.getLogger("TrainController") diff --git a/areal/core/dist_rollout.py b/areal/infra/dist_rollout.py similarity index 99% rename from areal/core/dist_rollout.py rename to areal/infra/dist_rollout.py index c09181312d..5800fc5439 100644 --- a/areal/core/dist_rollout.py +++ b/areal/infra/dist_rollout.py @@ -8,7 +8,7 @@ from areal.api.engine_api import InferenceEngine, TrainEngine from areal.api.workflow_api import WorkflowLike -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.data import ( all_gather_tensor_container, broadcast_tensor_container, diff --git a/areal/platforms/__init__.py b/areal/infra/platforms/__init__.py similarity index 100% rename from areal/platforms/__init__.py rename to areal/infra/platforms/__init__.py diff --git a/areal/platforms/cpu.py b/areal/infra/platforms/cpu.py similarity index 100% rename from areal/platforms/cpu.py rename to areal/infra/platforms/cpu.py diff --git a/areal/platforms/cuda.py b/areal/infra/platforms/cuda.py similarity index 100% rename from areal/platforms/cuda.py rename to areal/infra/platforms/cuda.py diff --git a/areal/platforms/npu.py b/areal/infra/platforms/npu.py similarity index 100% rename from areal/platforms/npu.py rename to areal/infra/platforms/npu.py diff --git a/areal/platforms/platform.py b/areal/infra/platforms/platform.py similarity index 100% rename from areal/platforms/platform.py rename to areal/infra/platforms/platform.py diff --git a/areal/platforms/unknown.py b/areal/infra/platforms/unknown.py similarity index 100% rename from areal/platforms/unknown.py rename to areal/infra/platforms/unknown.py diff --git a/areal/core/remote_inf_engine.py b/areal/infra/remote_inf_engine.py similarity index 99% rename from areal/core/remote_inf_engine.py rename to areal/infra/remote_inf_engine.py index f713a83a77..43822b2425 100644 --- a/areal/core/remote_inf_engine.py +++ b/areal/infra/remote_inf_engine.py @@ -34,8 +34,8 @@ WeightUpdateRequests, ) from areal.api.workflow_api import AgentWorkflow, RolloutWorkflow, WorkflowLike -from areal.core import workflow_context -from areal.platforms import current_platform +from areal.infra import workflow_context +from areal.infra.platforms import current_platform from areal.utils import logging, name_resolve, names from areal.utils.concurrent import get_executor from areal.utils.data import concat_padded_tensors diff --git a/areal/core/staleness_manager.py b/areal/infra/staleness_manager.py similarity index 100% rename from areal/core/staleness_manager.py rename to areal/infra/staleness_manager.py diff --git a/areal/core/workflow_context.py b/areal/infra/workflow_context.py similarity index 100% rename from areal/core/workflow_context.py rename to areal/infra/workflow_context.py diff --git a/areal/core/workflow_executor.py b/areal/infra/workflow_executor.py similarity index 99% rename from areal/core/workflow_executor.py rename to areal/infra/workflow_executor.py index e37c1804e1..01f57889de 100644 --- a/areal/core/workflow_executor.py +++ b/areal/infra/workflow_executor.py @@ -20,14 +20,14 @@ from areal.api.cli_args import InferenceEngineConfig from areal.api.workflow_api import RolloutWorkflow -from areal.core.async_task_runner import ( +from .async_task_runner import ( AsyncTaskRunner, TaskQueueFullError, TimedResult, ) -from areal.core.staleness_manager import StalenessManager -from areal.core import workflow_context -from areal.core.workflow_context import WorkflowContext +from .staleness_manager import StalenessManager +from areal.infra import workflow_context +from .workflow_context import WorkflowContext from areal.experimental.openai.types import InteractionWithTokenLogpReward from areal.utils import logging, perf_tracer, stats_tracker from areal.utils.concurrent import get_executor diff --git a/areal/launcher/local.py b/areal/launcher/local.py index 69aa389d90..01a83fbe81 100644 --- a/areal/launcher/local.py +++ b/areal/launcher/local.py @@ -18,7 +18,7 @@ to_structured_cfg, vLLMConfig, ) -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils import logging, name_resolve, names from areal.utils.exp_metadata import save_experiment_metadata from areal.utils.launcher import ( diff --git a/areal/launcher/ray.py b/areal/launcher/ray.py index fab9067863..10112839a4 100644 --- a/areal/launcher/ray.py +++ b/areal/launcher/ray.py @@ -22,7 +22,7 @@ to_structured_cfg, vLLMConfig, ) -from areal.platforms import current_platform, is_npu_available +from areal.infra.platforms import current_platform, is_npu_available from areal.utils import name_resolve, names from areal.utils.exp_metadata import save_experiment_metadata from areal.utils.launcher import ( diff --git a/areal/launcher/sglang_server.py b/areal/launcher/sglang_server.py index 1e4133b719..b66b97d58a 100644 --- a/areal/launcher/sglang_server.py +++ b/areal/launcher/sglang_server.py @@ -17,7 +17,7 @@ parse_cli_args, to_structured_cfg, ) -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils import logging, name_resolve, names from areal.utils.launcher import ( TRITON_CACHE_PATH, diff --git a/areal/launcher/slurm.py b/areal/launcher/slurm.py index 8ed22fa1cd..662f836bd1 100644 --- a/areal/launcher/slurm.py +++ b/areal/launcher/slurm.py @@ -16,7 +16,7 @@ to_structured_cfg, vLLMConfig, ) -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils import name_resolve, names from areal.utils.exp_metadata import save_experiment_metadata from areal.utils.launcher import ( diff --git a/areal/launcher/vllm_server.py b/areal/launcher/vllm_server.py index c9d3b279c4..54aa79f502 100644 --- a/areal/launcher/vllm_server.py +++ b/areal/launcher/vllm_server.py @@ -18,7 +18,7 @@ vLLMConfig, ) from areal.api.io_struct import AllocationMode -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils import logging, name_resolve, names from areal.utils.launcher import TRITON_CACHE_PATH, get_scheduling_spec from areal.utils.network import find_free_ports, gethostip diff --git a/areal/models/mcore/hf_load.py b/areal/models/mcore/hf_load.py index 6510bc3648..089fdb7f46 100644 --- a/areal/models/mcore/hf_load.py +++ b/areal/models/mcore/hf_load.py @@ -12,8 +12,8 @@ from megatron.core.fp8_utils import is_float8tensor from safetensors import safe_open +from areal.infra.platforms import current_platform from areal.models.mcore.registry import unwrap_to_gpt_model -from areal.platforms import current_platform from areal.utils import logging from areal.utils.fp8 import ( FP8BlockwiseTensorHelper, diff --git a/areal/models/mcore/hf_save.py b/areal/models/mcore/hf_save.py index 2a629b765f..7111292c9b 100644 --- a/areal/models/mcore/hf_save.py +++ b/areal/models/mcore/hf_save.py @@ -14,8 +14,8 @@ from safetensors.torch import save_file from torch.distributed._functional_collectives import all_gather_into_tensor_coalesced +from areal.infra.platforms import current_platform from areal.models.mcore.registry import unwrap_to_gpt_model -from areal.platforms import current_platform from areal.utils import logging from areal.utils.fp8 import ( FP8BlockwiseTensorHelper, diff --git a/areal/scheduler/local.py b/areal/scheduler/local.py index d3325e7011..0266863396 100644 --- a/areal/scheduler/local.py +++ b/areal/scheduler/local.py @@ -19,7 +19,7 @@ SchedulingStrategyType, ) from areal.api.scheduler_api import Job, Scheduler, SchedulingSpec, Worker -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.scheduler.exceptions import ( EngineCallError, EngineCreationError, diff --git a/areal/scheduler/ray.py b/areal/scheduler/ray.py index 5182a208fa..4fedbaf6d8 100644 --- a/areal/scheduler/ray.py +++ b/areal/scheduler/ray.py @@ -41,7 +41,7 @@ def ray_resource_type(): if torch.cuda.is_available(): return "GPU" - from areal.platforms import is_npu_available + from areal.infra.platforms import is_npu_available if is_npu_available: return "NPU" diff --git a/areal/scheduler/rpc/ray_rpc_server.py b/areal/scheduler/rpc/ray_rpc_server.py index 74c0949e8b..59bc601a91 100644 --- a/areal/scheduler/rpc/ray_rpc_server.py +++ b/areal/scheduler/rpc/ray_rpc_server.py @@ -39,7 +39,7 @@ def __init__(self): def _get_device(self): # lazy resolve the device inside worker process - from areal.platforms import current_platform + from areal.infra.platforms import current_platform return current_platform.current_device() diff --git a/areal/scheduler/rpc/rpc_server.py b/areal/scheduler/rpc/rpc_server.py index 6ea7098eb4..55f97cad6c 100644 --- a/areal/scheduler/rpc/rpc_server.py +++ b/areal/scheduler/rpc/rpc_server.py @@ -20,7 +20,7 @@ from areal.api.cli_args import BaseExperimentConfig, NameResolveConfig from areal.api.engine_api import InferenceEngine, TrainEngine -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.scheduler.rpc import rtensor from areal.scheduler.rpc.rtensor import RTensor from areal.scheduler.rpc.serialization import ( diff --git a/areal/tests/experimental/archon/test_checkpoint_e2e.py b/areal/tests/experimental/archon/test_checkpoint_e2e.py index 6ae5d43ccc..dd6f37498c 100644 --- a/areal/tests/experimental/archon/test_checkpoint_e2e.py +++ b/areal/tests/experimental/archon/test_checkpoint_e2e.py @@ -24,7 +24,7 @@ import pytest import torch -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.tests.experimental.archon.utils import ( DENSE_MODEL_PATHS, MOE_MODEL_PATHS, diff --git a/areal/tests/experimental/archon/test_distributed_cp.py b/areal/tests/experimental/archon/test_distributed_cp.py index 12b4b8d630..3ebd882d1b 100644 --- a/areal/tests/experimental/archon/test_distributed_cp.py +++ b/areal/tests/experimental/archon/test_distributed_cp.py @@ -9,7 +9,7 @@ import pytest import torch -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.network import find_free_ports pytestmark = pytest.mark.skipif( diff --git a/areal/tests/experimental/archon/test_distributed_dp.py b/areal/tests/experimental/archon/test_distributed_dp.py index e3f695a38c..e674f28489 100644 --- a/areal/tests/experimental/archon/test_distributed_dp.py +++ b/areal/tests/experimental/archon/test_distributed_dp.py @@ -7,7 +7,7 @@ import pytest import torch -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.tests.experimental.archon.utils import run_torchrun_test pytestmark = pytest.mark.skipif( diff --git a/areal/tests/experimental/archon/test_distributed_ep.py b/areal/tests/experimental/archon/test_distributed_ep.py index 7330a5cbf3..1a42b3f784 100644 --- a/areal/tests/experimental/archon/test_distributed_ep.py +++ b/areal/tests/experimental/archon/test_distributed_ep.py @@ -11,7 +11,7 @@ import pytest import torch -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.network import find_free_ports pytestmark = pytest.mark.skipif( diff --git a/areal/tests/experimental/archon/test_distributed_etp.py b/areal/tests/experimental/archon/test_distributed_etp.py index 964fca6063..7fe05f873e 100644 --- a/areal/tests/experimental/archon/test_distributed_etp.py +++ b/areal/tests/experimental/archon/test_distributed_etp.py @@ -11,7 +11,7 @@ import pytest import torch -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.network import find_free_ports pytestmark = pytest.mark.skipif( diff --git a/areal/tests/experimental/archon/test_distributed_pp.py b/areal/tests/experimental/archon/test_distributed_pp.py index f86baa9953..2fc19e4207 100644 --- a/areal/tests/experimental/archon/test_distributed_pp.py +++ b/areal/tests/experimental/archon/test_distributed_pp.py @@ -32,7 +32,7 @@ import pytest import torch -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.network import find_free_ports pytestmark = pytest.mark.skipif( diff --git a/areal/tests/experimental/archon/test_distributed_tp.py b/areal/tests/experimental/archon/test_distributed_tp.py index ed2dc42086..0e84d97f32 100644 --- a/areal/tests/experimental/archon/test_distributed_tp.py +++ b/areal/tests/experimental/archon/test_distributed_tp.py @@ -9,7 +9,7 @@ import pytest import torch -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.network import find_free_ports pytestmark = pytest.mark.skipif( diff --git a/areal/tests/experimental/archon/test_forward.py b/areal/tests/experimental/archon/test_forward.py index fb30db2fbd..a4d9f81c7b 100644 --- a/areal/tests/experimental/archon/test_forward.py +++ b/areal/tests/experimental/archon/test_forward.py @@ -18,7 +18,7 @@ import torch.distributed as dist from transformers import AutoConfig -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.tests.experimental.archon.utils import ( MODEL_PATHS, compare_logprobs, diff --git a/areal/tests/experimental/archon/test_grpo.py b/areal/tests/experimental/archon/test_grpo.py index 8315d9fa90..189ace7072 100644 --- a/areal/tests/experimental/archon/test_grpo.py +++ b/areal/tests/experimental/archon/test_grpo.py @@ -13,7 +13,7 @@ import torch from areal.engine.ppo.actor import grpo_loss_fn -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.tests.experimental.archon.utils import ( ComparisonMetrics, DualEngineFixture, diff --git a/areal/tests/experimental/archon/test_parallel_dims.py b/areal/tests/experimental/archon/test_parallel_dims.py index a5d0d5787b..71796c109b 100644 --- a/areal/tests/experimental/archon/test_parallel_dims.py +++ b/areal/tests/experimental/archon/test_parallel_dims.py @@ -13,7 +13,7 @@ import torch from areal.experimental.models.archon import ArchonParallelDims -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.network import find_free_ports diff --git a/areal/tests/experimental/archon/test_qwen3_parallelize.py b/areal/tests/experimental/archon/test_qwen3_parallelize.py index 505ebd735c..86e29f4743 100644 --- a/areal/tests/experimental/archon/test_qwen3_parallelize.py +++ b/areal/tests/experimental/archon/test_qwen3_parallelize.py @@ -20,7 +20,7 @@ import pytest import torch -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.network import find_free_ports pytestmark = pytest.mark.skipif( diff --git a/areal/tests/experimental/archon/test_varlen_attention.py b/areal/tests/experimental/archon/test_varlen_attention.py index 916413bbed..21e50efb14 100644 --- a/areal/tests/experimental/archon/test_varlen_attention.py +++ b/areal/tests/experimental/archon/test_varlen_attention.py @@ -3,7 +3,7 @@ import pytest import torch -from areal.platforms import current_platform +from areal.infra.platforms import current_platform def setup_environment(): diff --git a/areal/tests/experimental/archon/torchrun/run_cp_forward.py b/areal/tests/experimental/archon/torchrun/run_cp_forward.py index 1a0898ea12..bf46918d27 100644 --- a/areal/tests/experimental/archon/torchrun/run_cp_forward.py +++ b/areal/tests/experimental/archon/torchrun/run_cp_forward.py @@ -21,7 +21,7 @@ ) from areal.api.io_struct import FinetuneSpec from areal.experimental.engine.archon_engine import ArchonEngine -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.tests.utils import get_model_path from areal.utils.data import tensor_container_to diff --git a/areal/tests/experimental/archon/torchrun/run_forward.py b/areal/tests/experimental/archon/torchrun/run_forward.py index d393a4962d..5e9476479f 100644 --- a/areal/tests/experimental/archon/torchrun/run_forward.py +++ b/areal/tests/experimental/archon/torchrun/run_forward.py @@ -21,7 +21,7 @@ ) from areal.api.io_struct import FinetuneSpec from areal.experimental.engine.archon_engine import ArchonEngine -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.tests.utils import get_model_path from areal.utils.data import tensor_container_to diff --git a/areal/tests/experimental/archon/torchrun/run_tp_forward.py b/areal/tests/experimental/archon/torchrun/run_tp_forward.py index c28369d17b..e8200f7024 100644 --- a/areal/tests/experimental/archon/torchrun/run_tp_forward.py +++ b/areal/tests/experimental/archon/torchrun/run_tp_forward.py @@ -21,7 +21,7 @@ ) from areal.api.io_struct import FinetuneSpec from areal.experimental.engine.archon_engine import ArchonEngine -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.tests.utils import get_model_path from areal.utils.data import tensor_container_to diff --git a/areal/tests/experimental/archon/torchrun/run_vs_fsdp.py b/areal/tests/experimental/archon/torchrun/run_vs_fsdp.py index b103734642..b33c5f3244 100644 --- a/areal/tests/experimental/archon/torchrun/run_vs_fsdp.py +++ b/areal/tests/experimental/archon/torchrun/run_vs_fsdp.py @@ -19,7 +19,7 @@ from areal.api.io_struct import FinetuneSpec from areal.engine.fsdp_engine import FSDPLMEngine from areal.experimental.engine.archon_engine import ArchonLMEngine -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.tests.experimental.archon.torchrun.dist_utils import write_result from areal.tests.utils import get_dataset_path, get_model_path from areal.utils.data import pad_sequences_to_tensors diff --git a/areal/tests/experimental/archon/utils.py b/areal/tests/experimental/archon/utils.py index 77f989aa97..6229cca523 100644 --- a/areal/tests/experimental/archon/utils.py +++ b/areal/tests/experimental/archon/utils.py @@ -17,7 +17,7 @@ from areal.engine.fsdp_engine import FSDPLMEngine from areal.experimental.engine.archon_engine import ArchonLMEngine from areal.experimental.models.archon import get_model_spec, is_supported_model -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.tests.utils import get_dataset_path, get_model_path from areal.utils.data import pad_sequences_to_tensors from areal.utils.hf_utils import load_hf_tokenizer diff --git a/areal/tests/experimental/openai/test_proxy_integration.py b/areal/tests/experimental/openai/test_proxy_integration.py index 61b206ea70..67b84f8f8b 100644 --- a/areal/tests/experimental/openai/test_proxy_integration.py +++ b/areal/tests/experimental/openai/test_proxy_integration.py @@ -15,8 +15,8 @@ SGLangConfig, ) from areal.api.io_struct import LocalInfServerInfo -from areal.controller.rollout_controller import RolloutController from areal.engine.sglang_remote import RemoteSGLangEngine +from areal.infra import RolloutController from areal.scheduler.local import LocalScheduler from areal.scheduler.rpc.rtensor import RTensor from areal.tests.utils import get_model_path diff --git a/areal/tests/fp8/test_fp8_bf16_comparison.py b/areal/tests/fp8/test_fp8_bf16_comparison.py index 477da979ab..ab98be1141 100644 --- a/areal/tests/fp8/test_fp8_bf16_comparison.py +++ b/areal/tests/fp8/test_fp8_bf16_comparison.py @@ -17,7 +17,7 @@ import torch.nn.functional as F from transformers import AutoTokenizer -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.tests.fp8.comparison_utils import ( compare_logits, compare_tensors_dict, diff --git a/areal/tests/fp8/test_fp8_conversion.py b/areal/tests/fp8/test_fp8_conversion.py index 9537e2be13..d592fae0f0 100644 --- a/areal/tests/fp8/test_fp8_conversion.py +++ b/areal/tests/fp8/test_fp8_conversion.py @@ -8,7 +8,7 @@ import pytest import torch -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils import logging logger = logging.getLogger("Test FP8 Conversion") diff --git a/areal/tests/test_async_task_runner.py b/areal/tests/test_async_task_runner.py index 4ae0643021..8ff0ec9159 100644 --- a/areal/tests/test_async_task_runner.py +++ b/areal/tests/test_async_task_runner.py @@ -10,7 +10,7 @@ import pytest -from areal.core.async_task_runner import AsyncTaskRunner +from areal.infra.async_task_runner import AsyncTaskRunner class TestAsyncTaskRunnerBasic: diff --git a/areal/tests/test_data_redistribution.py b/areal/tests/test_data_redistribution.py index 8922e3b7d5..f4baa4a487 100644 --- a/areal/tests/test_data_redistribution.py +++ b/areal/tests/test_data_redistribution.py @@ -6,7 +6,7 @@ import torch from torch.testing import assert_close -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.data import concat_padded_tensors from areal.utils.environ import is_in_ci from areal.utils.network import find_free_ports diff --git a/areal/tests/test_examples.py b/areal/tests/test_examples.py index 0b12cbbfda..32a1af86ef 100644 --- a/areal/tests/test_examples.py +++ b/areal/tests/test_examples.py @@ -10,7 +10,7 @@ import pytest -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.tests.utils import get_dataset_path, get_model_path from areal.utils import logging from areal.utils.concurrent import run_async_task diff --git a/areal/tests/test_fsdp_dcp.py b/areal/tests/test_fsdp_dcp.py index 51fa0ada2f..765911b6b6 100644 --- a/areal/tests/test_fsdp_dcp.py +++ b/areal/tests/test_fsdp_dcp.py @@ -3,7 +3,7 @@ import pytest from areal.api.alloc_mode import AllocationMode -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.network import find_free_ports diff --git a/areal/tests/test_fsdp_memory_efficient_lora.py b/areal/tests/test_fsdp_memory_efficient_lora.py index 63b6774d38..c98e952334 100644 --- a/areal/tests/test_fsdp_memory_efficient_lora.py +++ b/areal/tests/test_fsdp_memory_efficient_lora.py @@ -3,7 +3,7 @@ import pytest from areal.api.alloc_mode import AllocationMode -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.network import find_free_ports diff --git a/areal/tests/test_fsdp_ulysses_forward.py b/areal/tests/test_fsdp_ulysses_forward.py index 28df2ba7c9..a9f8ee0fea 100644 --- a/areal/tests/test_fsdp_ulysses_forward.py +++ b/areal/tests/test_fsdp_ulysses_forward.py @@ -2,7 +2,7 @@ import pytest -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.network import find_free_ports diff --git a/areal/tests/test_fsdp_ulysses_train_batch.py b/areal/tests/test_fsdp_ulysses_train_batch.py index 08ddc4d609..67dedf720e 100644 --- a/areal/tests/test_fsdp_ulysses_train_batch.py +++ b/areal/tests/test_fsdp_ulysses_train_batch.py @@ -2,7 +2,7 @@ import pytest -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.network import find_free_ports diff --git a/areal/tests/test_megatron_engine.py b/areal/tests/test_megatron_engine.py index ab698013e3..2e73b8ed1e 100644 --- a/areal/tests/test_megatron_engine.py +++ b/areal/tests/test_megatron_engine.py @@ -16,7 +16,7 @@ ) from areal.api.io_struct import FinetuneSpec, SaveLoadMeta from areal.engine.megatron_engine import MegatronEngine -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.tests.utils import get_model_path from areal.utils import logging diff --git a/areal/tests/test_megatron_engine_distributed.py b/areal/tests/test_megatron_engine_distributed.py index 151c287949..bb4749b26b 100644 --- a/areal/tests/test_megatron_engine_distributed.py +++ b/areal/tests/test_megatron_engine_distributed.py @@ -3,7 +3,7 @@ import pytest from areal.api.alloc_mode import AllocationMode -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.network import find_free_ports diff --git a/areal/tests/test_offload.py b/areal/tests/test_offload.py index 31e4877301..0173972b28 100644 --- a/areal/tests/test_offload.py +++ b/areal/tests/test_offload.py @@ -14,7 +14,7 @@ from areal.api.io_struct import FinetuneSpec from areal.engine.fsdp_engine import FSDPEngine from areal.engine.megatron_engine import MegatronEngine -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.tests.utils import get_model_path from areal.utils.network import find_free_ports from areal.utils.offload import get_tms_env_vars diff --git a/areal/tests/test_packed_vs_padded_consistency.py b/areal/tests/test_packed_vs_padded_consistency.py index 46cdbe0096..f253f8c068 100644 --- a/areal/tests/test_packed_vs_padded_consistency.py +++ b/areal/tests/test_packed_vs_padded_consistency.py @@ -8,7 +8,7 @@ from areal.api.cli_args import TrainEngineConfig from areal.engine.fsdp_engine import FSDPEngine -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.tests.utils import get_model_path from areal.utils.data import concat_padded_tensors, tensor_container_to from areal.utils.hf_utils import load_hf_processor_and_tokenizer diff --git a/areal/tests/test_perf_tracer.py b/areal/tests/test_perf_tracer.py index f56a5cba19..4989384ba6 100644 --- a/areal/tests/test_perf_tracer.py +++ b/areal/tests/test_perf_tracer.py @@ -9,7 +9,7 @@ import pytest from areal.api.cli_args import PerfTracerConfig, SessionTracerConfig -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils import perf_tracer from areal.utils.network import find_free_ports from areal.utils.perf_tracer import Category, SessionTraceEvent diff --git a/areal/tests/test_rollout_controller.py b/areal/tests/test_rollout_controller.py index 9f4099c2ec..32c358a1bb 100644 --- a/areal/tests/test_rollout_controller.py +++ b/areal/tests/test_rollout_controller.py @@ -14,8 +14,8 @@ ) from areal.api.io_struct import ModelRequest, ParamSpec, WeightUpdateMeta from areal.api.scheduler_api import Worker -from areal.controller import RolloutController from areal.engine.sglang_remote import RemoteSGLangEngine +from areal.infra import RolloutController from areal.scheduler.local import LocalScheduler from areal.tests.utils import get_model_path from areal.utils.hf_utils import load_hf_tokenizer diff --git a/areal/tests/test_staleness_manager.py b/areal/tests/test_staleness_manager.py index 13b1f4954a..67d6a2bb6c 100644 --- a/areal/tests/test_staleness_manager.py +++ b/areal/tests/test_staleness_manager.py @@ -9,7 +9,7 @@ import pytest -from areal.core.staleness_manager import StalenessManager +from areal.infra.staleness_manager import StalenessManager class MockVersionProvider: diff --git a/areal/tests/test_train_controller.py b/areal/tests/test_train_controller.py index 72d844052d..79f02b7c69 100644 --- a/areal/tests/test_train_controller.py +++ b/areal/tests/test_train_controller.py @@ -20,7 +20,7 @@ WeightUpdateMeta, ) from areal.api.scheduler_api import Worker -from areal.controller.train_controller import TrainController +from areal.infra import TrainController class MockTrainEngine(TrainEngine): diff --git a/areal/tests/test_train_engine.py b/areal/tests/test_train_engine.py index af76761ff9..80cc35c312 100644 --- a/areal/tests/test_train_engine.py +++ b/areal/tests/test_train_engine.py @@ -10,7 +10,7 @@ from areal.api.cli_args import MicroBatchSpec, OptimizerConfig, TrainEngineConfig from areal.api.io_struct import FinetuneSpec, SaveLoadMeta -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.tests.utils import get_model_path VOCAB_SIZE = 100 diff --git a/areal/tests/test_tree_training.py b/areal/tests/test_tree_training.py index 85f79bdb63..3c8a5fd263 100644 --- a/areal/tests/test_tree_training.py +++ b/areal/tests/test_tree_training.py @@ -12,8 +12,8 @@ from areal.api.io_struct import FinetuneSpec from areal.engine.fsdp_engine import FSDPEngine from areal.engine.megatron_engine import MegatronEngine +from areal.infra.platforms import current_platform from areal.models.tree_attn.module import restore_patch_fsdp_for_tree_training -from areal.platforms import current_platform from areal.tests.utils import get_model_path from areal.utils import logging diff --git a/areal/tests/test_ulysses.py b/areal/tests/test_ulysses.py index d3f4ce7d7a..cab512fd93 100644 --- a/areal/tests/test_ulysses.py +++ b/areal/tests/test_ulysses.py @@ -2,7 +2,7 @@ import pytest -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.network import find_free_ports diff --git a/areal/tests/test_ulysses_all_to_all.py b/areal/tests/test_ulysses_all_to_all.py index e17dac7ab4..c976f20f44 100644 --- a/areal/tests/test_ulysses_all_to_all.py +++ b/areal/tests/test_ulysses_all_to_all.py @@ -10,7 +10,7 @@ import pytest -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.network import find_free_ports diff --git a/areal/tests/test_vocab_parallel.py b/areal/tests/test_vocab_parallel.py index 28e51eb903..4996526340 100644 --- a/areal/tests/test_vocab_parallel.py +++ b/areal/tests/test_vocab_parallel.py @@ -2,7 +2,7 @@ import pytest -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.network import find_free_ports diff --git a/areal/tests/torchrun/redistribute.py b/areal/tests/torchrun/redistribute.py index 9237534c5c..6fca229f49 100644 --- a/areal/tests/torchrun/redistribute.py +++ b/areal/tests/torchrun/redistribute.py @@ -6,8 +6,8 @@ import torch import torch.distributed as dist -from areal.core.dist_rollout import redistribute_trajectories -from areal.platforms import current_platform +from areal.infra.dist_rollout import redistribute_trajectories +from areal.infra.platforms import current_platform from areal.utils.data import tensor_container_to diff --git a/areal/tests/torchrun/run_fsdp_dcp_distributed.py b/areal/tests/torchrun/run_fsdp_dcp_distributed.py index 8eae9c4859..c3295aa8de 100644 --- a/areal/tests/torchrun/run_fsdp_dcp_distributed.py +++ b/areal/tests/torchrun/run_fsdp_dcp_distributed.py @@ -13,7 +13,7 @@ from areal.api.cli_args import MicroBatchSpec, OptimizerConfig, TrainEngineConfig from areal.api.io_struct import FinetuneSpec, SaveLoadMeta from areal.engine.fsdp_engine import FSDPEngine -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.tests.utils import get_model_path from areal.utils import seeding diff --git a/areal/tests/torchrun/run_fsdp_memory_efficient_lora.py b/areal/tests/torchrun/run_fsdp_memory_efficient_lora.py index 0359a78060..a23a19475b 100644 --- a/areal/tests/torchrun/run_fsdp_memory_efficient_lora.py +++ b/areal/tests/torchrun/run_fsdp_memory_efficient_lora.py @@ -15,7 +15,7 @@ ) from areal.api.io_struct import FinetuneSpec from areal.engine.fsdp_engine import FSDPEngine -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.tests.utils import get_model_path MODEL_PATH = get_model_path( diff --git a/areal/tests/torchrun/run_fsdp_ulysses_forward.py b/areal/tests/torchrun/run_fsdp_ulysses_forward.py index 8fac682a33..9797149108 100644 --- a/areal/tests/torchrun/run_fsdp_ulysses_forward.py +++ b/areal/tests/torchrun/run_fsdp_ulysses_forward.py @@ -13,7 +13,7 @@ ) from areal.api.io_struct import FinetuneSpec from areal.engine.fsdp_engine import FSDPEngine -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.tests.utils import get_model_path MODEL_PATHS = { diff --git a/areal/tests/torchrun/run_fsdp_ulysses_train_batch.py b/areal/tests/torchrun/run_fsdp_ulysses_train_batch.py index 29ff607b64..968f8e2b96 100644 --- a/areal/tests/torchrun/run_fsdp_ulysses_train_batch.py +++ b/areal/tests/torchrun/run_fsdp_ulysses_train_batch.py @@ -13,7 +13,7 @@ ) from areal.api.io_struct import FinetuneSpec from areal.engine.fsdp_engine import FSDPEngine -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.tests.utils import get_model_path from areal.utils.data import tensor_container_to diff --git a/areal/tests/torchrun/run_lock.py b/areal/tests/torchrun/run_lock.py index 4da817fc24..287909b7e9 100644 --- a/areal/tests/torchrun/run_lock.py +++ b/areal/tests/torchrun/run_lock.py @@ -6,7 +6,7 @@ import torch.distributed as dist -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.lock import DistributedLock diff --git a/areal/tests/torchrun/run_megatron_engine_distributed.py b/areal/tests/torchrun/run_megatron_engine_distributed.py index cfe33c9275..988fc0763d 100644 --- a/areal/tests/torchrun/run_megatron_engine_distributed.py +++ b/areal/tests/torchrun/run_megatron_engine_distributed.py @@ -19,7 +19,7 @@ from areal.api.io_struct import FinetuneSpec, SaveLoadMeta from areal.engine.fsdp_engine import FSDPEngine from areal.engine.megatron_engine import MegatronEngine -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.tests.utils import get_model_path from areal.utils import seeding from areal.utils.data import broadcast_tensor_container diff --git a/areal/tests/torchrun/run_ulysses.py b/areal/tests/torchrun/run_ulysses.py index 56b307f24d..3b6044f75b 100644 --- a/areal/tests/torchrun/run_ulysses.py +++ b/areal/tests/torchrun/run_ulysses.py @@ -8,9 +8,9 @@ Qwen2RotaryEmbedding, ) +from areal.infra.platforms import current_platform from areal.models.fsdp.ulysses import set_ulysses_sequence_parallel_group from areal.models.transformers.ulyssess_patch import apply_monkey_patch -from areal.platforms import current_platform from areal.utils.constants import DIST_GROUP_DEFAULT_TIMEOUT diff --git a/areal/tests/torchrun/run_ulysses_all_to_all.py b/areal/tests/torchrun/run_ulysses_all_to_all.py index 933cfaf674..c9eaf54748 100644 --- a/areal/tests/torchrun/run_ulysses_all_to_all.py +++ b/areal/tests/torchrun/run_ulysses_all_to_all.py @@ -12,11 +12,11 @@ import torch import torch.distributed as dist +from areal.infra.platforms import current_platform from areal.models.fsdp.ulysses import ( all_to_all_tensor, set_ulysses_sequence_parallel_group, ) -from areal.platforms import current_platform from areal.utils.constants import DIST_GROUP_DEFAULT_TIMEOUT diff --git a/areal/tests/torchrun/run_vocab_parallel.py b/areal/tests/torchrun/run_vocab_parallel.py index 2f7aa56e3c..769be34dcd 100644 --- a/areal/tests/torchrun/run_vocab_parallel.py +++ b/areal/tests/torchrun/run_vocab_parallel.py @@ -3,7 +3,7 @@ import torch import torch.distributed as dist -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils.functional.vocab_parallel import ( _vocab_parallel_logprobs, _vocab_parallel_logprobs_entropy, diff --git a/areal/tools/profile_archon.py b/areal/tools/profile_archon.py index 42a148f547..52f8f2d47d 100644 --- a/areal/tools/profile_archon.py +++ b/areal/tools/profile_archon.py @@ -36,7 +36,7 @@ import torch from torch.profiler import ProfilerActivity -from areal.platforms import current_platform +from areal.infra import current_platform def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: diff --git a/areal/tools/profile_engines.py b/areal/tools/profile_engines.py index f5c3be9cee..62dfcd72ee 100644 --- a/areal/tools/profile_engines.py +++ b/areal/tools/profile_engines.py @@ -31,7 +31,7 @@ import torch from torch.profiler import ProfilerActivity -from areal.platforms import current_platform +from areal.infra import current_platform @dataclass diff --git a/areal/tools/profile_fsdp.py b/areal/tools/profile_fsdp.py index 533e4292a4..c4f4261d85 100644 --- a/areal/tools/profile_fsdp.py +++ b/areal/tools/profile_fsdp.py @@ -39,7 +39,7 @@ from torch.profiler import ProfilerActivity from transformers import AutoConfig, AutoModelForCausalLM -from areal.platforms import current_platform +from areal.infra import current_platform def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: diff --git a/areal/utils/data.py b/areal/utils/data.py index 73eb28d1e5..62b991f2cd 100644 --- a/areal/utils/data.py +++ b/areal/utils/data.py @@ -14,7 +14,7 @@ from torchdata.stateful_dataloader import StatefulDataLoader from areal.api.cli_args import MicroBatchSpec, NormConfig -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils import datapack, logging from areal.utils.math import align diff --git a/areal/utils/fp8/tensor_helper.py b/areal/utils/fp8/tensor_helper.py index ebd66a821a..c6abece4bf 100644 --- a/areal/utils/fp8/tensor_helper.py +++ b/areal/utils/fp8/tensor_helper.py @@ -3,7 +3,7 @@ import torch from megatron.core.fp8_utils import is_float8tensor -from areal.platforms import current_platform +from areal.infra.platforms import current_platform class FP8BlockwiseTensorHelper(torch.Tensor): diff --git a/areal/utils/fsdp/__init__.py b/areal/utils/fsdp/__init__.py index 2f3b152ff4..655cf46409 100644 --- a/areal/utils/fsdp/__init__.py +++ b/areal/utils/fsdp/__init__.py @@ -11,7 +11,7 @@ ) from transformers import PreTrainedModel -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils import logging, pkg_version if pkg_version.is_version_greater_or_equal("torch", "2.6.0"): diff --git a/areal/utils/fsdp/grad.py b/areal/utils/fsdp/grad.py index a433bf2a4f..48fd5c6e7c 100644 --- a/areal/utils/fsdp/grad.py +++ b/areal/utils/fsdp/grad.py @@ -7,7 +7,7 @@ from torch.distributed.device_mesh import DeviceMesh from torch.distributed.tensor import DTensor -from areal.platforms import current_platform +from areal.infra.platforms import current_platform __all__ = [ "fsdp2_clip_grad_norm", diff --git a/areal/utils/fsdp/multi_tensor_apply.py b/areal/utils/fsdp/multi_tensor_apply.py index bb8dd64664..a2c30ae8bf 100644 --- a/areal/utils/fsdp/multi_tensor_apply.py +++ b/areal/utils/fsdp/multi_tensor_apply.py @@ -2,7 +2,7 @@ import torch -from areal.platforms import current_platform +from areal.infra.platforms import current_platform def local_multi_tensor_applier(op, noop_flag_buffer, tensor_lists, *args): diff --git a/areal/utils/fsdp/parallel.py b/areal/utils/fsdp/parallel.py index a5394f24bd..7ad7f7273e 100644 --- a/areal/utils/fsdp/parallel.py +++ b/areal/utils/fsdp/parallel.py @@ -18,8 +18,8 @@ from areal.api.alloc_mode import FSDPParallelStrategy from areal.api.cli_args import FSDPWrapPolicy, TrainEngineConfig +from areal.infra.platforms import current_platform from areal.models.parallel_styles import ReplicateParallel -from areal.platforms import current_platform from areal.utils.fsdp import apply_fsdp2 from areal.utils.model import ( is_gemma3_model, diff --git a/areal/utils/functional/vocab_parallel.py b/areal/utils/functional/vocab_parallel.py index fe7d57c499..e1fa00eae5 100644 --- a/areal/utils/functional/vocab_parallel.py +++ b/areal/utils/functional/vocab_parallel.py @@ -4,7 +4,7 @@ import torch from torch import distributed as dist -from areal.platforms import is_npu_available +from areal.infra.platforms import is_npu_available def _gather_logprobs( diff --git a/areal/utils/megatron_checkpointer.py b/areal/utils/megatron_checkpointer.py index aaee935dca..1d6079c621 100644 --- a/areal/utils/megatron_checkpointer.py +++ b/areal/utils/megatron_checkpointer.py @@ -29,7 +29,7 @@ FullyParallelSaveStrategyWrapper, ) -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils import logging logger = logging.getLogger("MegatronCheckpointer") diff --git a/areal/utils/recover.py b/areal/utils/recover.py index f4e51a0f60..e4de736608 100644 --- a/areal/utils/recover.py +++ b/areal/utils/recover.py @@ -11,7 +11,7 @@ from areal.api.cli_args import RecoverConfig from areal.api.engine_api import InferenceEngine, TrainEngine from areal.api.io_struct import FinetuneSpec, SaveLoadMeta, StepInfo, WeightUpdateMeta -from areal.controller.train_controller import TrainController +from areal.infra import TrainController from areal.utils import logging, timeutil from areal.utils.evaluator import Evaluator from areal.utils.saver import Saver diff --git a/areal/utils/saver.py b/areal/utils/saver.py index 02db552573..dacdca321b 100644 --- a/areal/utils/saver.py +++ b/areal/utils/saver.py @@ -6,7 +6,7 @@ from areal.api.cli_args import SaverConfig from areal.api.engine_api import TrainEngine from areal.api.io_struct import FinetuneSpec, SaveLoadMeta -from areal.controller.train_controller import TrainController +from areal.infra import TrainController from areal.utils import timeutil diff --git a/areal/utils/seeding.py b/areal/utils/seeding.py index 955cfa7eea..1105826e72 100644 --- a/areal/utils/seeding.py +++ b/areal/utils/seeding.py @@ -6,7 +6,7 @@ import torch import transformers -from areal.platforms import current_platform +from areal.infra.platforms import current_platform _SEED = None _BASE_SEED = None diff --git a/areal/utils/stats_tracker.py b/areal/utils/stats_tracker.py index 2efb6e5790..384e0d78a9 100644 --- a/areal/utils/stats_tracker.py +++ b/areal/utils/stats_tracker.py @@ -7,7 +7,7 @@ import torch import torch.distributed as dist -from areal.platforms import current_platform +from areal.infra.platforms import current_platform from areal.utils import logging from areal.utils.datapack import flat2d diff --git a/areal/utils/timeutil.py b/areal/utils/timeutil.py index 3be242fd8d..2f7ffb8b01 100644 --- a/areal/utils/timeutil.py +++ b/areal/utils/timeutil.py @@ -7,7 +7,7 @@ import torch import torch.distributed as dist -from areal.platforms import current_platform +from areal.infra.platforms import current_platform INFINITE_DURATION = 60 * 60 * 24 * 365 * 1000 diff --git a/areal/workflow/multi_turn.py b/areal/workflow/multi_turn.py index 2987e6fec6..f24efc22b5 100644 --- a/areal/workflow/multi_turn.py +++ b/areal/workflow/multi_turn.py @@ -5,12 +5,12 @@ import torch from transformers import PreTrainedTokenizerFast +from areal import workflow_context from areal.api.cli_args import GenerationHyperparameters from areal.api.engine_api import InferenceEngine from areal.api.io_struct import ModelRequest from areal.api.reward_api import AsyncRewardWrapper from areal.api.workflow_api import RolloutWorkflow -from areal.core import workflow_context from areal.utils import logging, stats_tracker logger = logging.getLogger("MultiTurnWorkflow") diff --git a/areal/workflow/rlvr.py b/areal/workflow/rlvr.py index 8de3f62cbc..e8386813b2 100644 --- a/areal/workflow/rlvr.py +++ b/areal/workflow/rlvr.py @@ -5,12 +5,12 @@ import torch from transformers import PreTrainedTokenizerFast +from areal import workflow_context from areal.api.cli_args import GenerationHyperparameters from areal.api.engine_api import InferenceEngine from areal.api.io_struct import ModelRequest, ModelResponse from areal.api.reward_api import AsyncRewardWrapper from areal.api.workflow_api import RolloutWorkflow -from areal.core import workflow_context from areal.utils import logging, stats_tracker from areal.utils.dynamic_import import import_from_string from areal.utils.perf_tracer import ( diff --git a/areal/workflow/vision_rlvr.py b/areal/workflow/vision_rlvr.py index 1df0da7704..cbb2648f58 100644 --- a/areal/workflow/vision_rlvr.py +++ b/areal/workflow/vision_rlvr.py @@ -5,11 +5,11 @@ import torch from transformers import AutoProcessor, PreTrainedTokenizerFast +from areal import workflow_context from areal.api.cli_args import GenerationHyperparameters from areal.api.engine_api import InferenceEngine from areal.api.io_struct import ModelRequest, ModelResponse from areal.api.reward_api import AsyncRewardWrapper -from areal.core import workflow_context from areal.utils import logging, stats_tracker from areal.utils.dynamic_import import import_from_string from areal.utils.image import image2base64 diff --git a/docs/best_practices/perf_profiling.md b/docs/best_practices/perf_profiling.md index 1bd423af4c..ff425cbd24 100644 --- a/docs/best_practices/perf_profiling.md +++ b/docs/best_practices/perf_profiling.md @@ -136,7 +136,7 @@ does a single prompt take from submission to reward calculation?). **Example** (from `areal/workflow/rlvr.py`): ```python -from areal.core import workflow_context +from areal import workflow_context from areal.utils.perf_tracer import ( atrace_session_phase, session_context, @@ -283,7 +283,7 @@ snapshot"). - Creates a point-in-time marker (not a duration) - Useful for events that don't have a meaningful duration -**Example** (from `areal/core/workflow_executor.py`): +**Example** (from `areal/infra/workflow_executor.py`): ```python perf_tracer.instant( @@ -316,7 +316,7 @@ consumption). - Parameters: Use `session_id=` to target a specific session, or `task_id=` to target all sessions in a task -**Example** (from `areal/core/workflow_executor.py`): +**Example** (from `areal/infra/workflow_executor.py`): ```python from areal.utils.perf_tracer import trace_session_event diff --git a/docs/cli_reference.md b/docs/cli_reference.md index ce30433a22..9a5bef255e 100644 --- a/docs/cli_reference.md +++ b/docs/cli_reference.md @@ -601,7 +601,7 @@ https://docs.vllm.ai/en/stable/api/index.html for detailed documentation. | `no_enable_chunked_prefill` | boolean | `False` | - | | `no_enable_prefix_caching` | boolean | `True` | - | | `gpu_memory_utilization` | float | `0.9` | - | -| `worker_extension_cls` | string | `"areal.thirdparty.vllm.vllm_worker_extension.VLLMWorkerExtension"` | - | +| `worker_extension_cls` | string | `"areal.engine.vllm_ext.vllm_worker_extension.VLLMWorkerExtension"` | - | | `enable_sleep_mode` | boolean | `False` | - | | `uvicorn_log_level` | string | `"warning"` | - | | `enable_lora` | boolean | `False` | - | diff --git a/examples/alignment/hhrlhf_rw.py b/examples/alignment/hhrlhf_rw.py index 6f31b6c20f..09d0d27d48 100644 --- a/examples/alignment/hhrlhf_rw.py +++ b/examples/alignment/hhrlhf_rw.py @@ -3,12 +3,12 @@ import torch.distributed as dist +from areal import current_platform from areal.api.alloc_mode import AllocationMode from areal.api.cli_args import RWConfig, load_expr_config from areal.api.io_struct import FinetuneSpec, StepInfo from areal.dataset import get_custom_dataset from areal.engine.fsdp_engine import FSDPRWEngine -from areal.platforms import current_platform from areal.utils import seeding, stats_tracker from areal.utils.data import ( broadcast_tensor_container, diff --git a/examples/camel/train.py b/examples/camel/train.py index bb4ec9a8eb..e2e5f2b74e 100644 --- a/examples/camel/train.py +++ b/examples/camel/train.py @@ -3,10 +3,10 @@ from camel.agents import ChatAgent from transformers import PreTrainedTokenizerFast +from areal import workflow_context from areal.api.cli_args import GenerationHyperparameters, GRPOConfig, load_expr_config from areal.api.reward_api import AsyncRewardWrapper from areal.api.workflow_api import RolloutWorkflow -from areal.core import workflow_context from areal.dataset import get_custom_dataset from areal.experimental.camel.openai_model import AReaLOpenAICompatibleModel from areal.experimental.openai import ArealOpenAI diff --git a/examples/countdown/train.py b/examples/countdown/train.py index 1d95663523..e549a9b331 100644 --- a/examples/countdown/train.py +++ b/examples/countdown/train.py @@ -7,11 +7,11 @@ from reward_score import compute_score from transformers import PreTrainedTokenizerFast +from areal import workflow_context from areal.api.cli_args import GenerationHyperparameters, GRPOConfig, load_expr_config from areal.api.engine_api import InferenceEngine from areal.api.io_struct import ModelRequest from areal.api.workflow_api import RolloutWorkflow -from areal.core import workflow_context from areal.experimental.trainer import PPOTrainer from areal.utils import logging, stats_tracker diff --git a/examples/lora/gsm8k_grpo_lora.py b/examples/lora/gsm8k_grpo_lora.py index 64ae808c42..b30238cbb8 100644 --- a/examples/lora/gsm8k_grpo_lora.py +++ b/examples/lora/gsm8k_grpo_lora.py @@ -4,14 +4,14 @@ import torch.distributed as dist +from areal import current_platform from areal.api.alloc_mode import AllocationMode from areal.api.cli_args import GRPOConfig, load_expr_config from areal.api.io_struct import FinetuneSpec, StepInfo, WeightUpdateMeta -from areal.core.dist_rollout import redistribute_trajectories from areal.dataset import get_custom_dataset from areal.engine.fsdp_engine import FSDPPPOActor from areal.engine.sglang_remote import RemoteSGLangEngine -from areal.platforms import current_platform +from areal.infra.dist_rollout import redistribute_trajectories from areal.reward.gsm8k import gsm8k_reward_fn from areal.utils import seeding, stats_tracker from areal.utils.data import ( diff --git a/examples/lora/gsm8k_grpo_lora_vllm.py b/examples/lora/gsm8k_grpo_lora_vllm.py index 906224ba7f..b4aa5ebb53 100644 --- a/examples/lora/gsm8k_grpo_lora_vllm.py +++ b/examples/lora/gsm8k_grpo_lora_vllm.py @@ -5,13 +5,13 @@ import torch.distributed as dist +from areal import current_platform from areal.api.alloc_mode import AllocationMode from areal.api.cli_args import GRPOConfig, load_expr_config from areal.api.io_struct import FinetuneSpec, StepInfo, WeightUpdateMeta from areal.dataset import get_custom_dataset from areal.engine.fsdp_engine import FSDPPPOActor from areal.engine.vllm_remote import RemotevLLMEngine -from areal.platforms import current_platform from areal.reward.gsm8k import gsm8k_reward_fn from areal.utils import seeding, stats_tracker from areal.utils.dataloader import create_dataloader diff --git a/examples/multi_turn_math/gsm8k_rl_mt.py b/examples/multi_turn_math/gsm8k_rl_mt.py index e56cdab61d..622f53d67d 100644 --- a/examples/multi_turn_math/gsm8k_rl_mt.py +++ b/examples/multi_turn_math/gsm8k_rl_mt.py @@ -5,10 +5,10 @@ from openai.types.chat import ChatCompletion from transformers import PreTrainedTokenizerFast +from areal import workflow_context from areal.api.cli_args import GenerationHyperparameters, GRPOConfig, load_expr_config from areal.api.reward_api import AsyncRewardWrapper from areal.api.workflow_api import RolloutWorkflow -from areal.core import workflow_context from areal.dataset import get_custom_dataset from areal.experimental.openai import ArealOpenAI from areal.experimental.trainer import PPOTrainer diff --git a/examples/openai_agents/train_agents.py b/examples/openai_agents/train_agents.py index 5c217d6c93..2ceea69998 100644 --- a/examples/openai_agents/train_agents.py +++ b/examples/openai_agents/train_agents.py @@ -5,10 +5,10 @@ from agents import Runner as OpenAIRunner from transformers import PreTrainedTokenizerFast +from areal import workflow_context from areal.api.cli_args import GenerationHyperparameters, GRPOConfig, load_expr_config from areal.api.reward_api import AsyncRewardWrapper from areal.api.workflow_api import RolloutWorkflow -from areal.core import workflow_context from areal.dataset import get_custom_dataset from areal.experimental.openai import ArealOpenAI from areal.experimental.trainer import PPOTrainer diff --git a/examples/search_agent/tongyi_deepresearch/train.py b/examples/search_agent/tongyi_deepresearch/train.py index 7932081c33..fd1a0fc1d3 100644 --- a/examples/search_agent/tongyi_deepresearch/train.py +++ b/examples/search_agent/tongyi_deepresearch/train.py @@ -8,6 +8,7 @@ from datasets import load_dataset from transformers import PreTrainedTokenizerFast +from areal import workflow_context from areal.api.cli_args import ( GenerationHyperparameters, GRPOConfig, @@ -15,7 +16,6 @@ load_expr_config, ) from areal.api.workflow_api import RolloutWorkflow -from areal.core import workflow_context from areal.engine.sglang_remote import RemoteSGLangEngine from areal.experimental.openai import ArealOpenAI from areal.experimental.trainer import PPOTrainer diff --git a/examples/tir/tir_workflow.py b/examples/tir/tir_workflow.py index 64a2234942..8183b1d785 100644 --- a/examples/tir/tir_workflow.py +++ b/examples/tir/tir_workflow.py @@ -7,6 +7,7 @@ import torch from transformers import PreTrainedTokenizerFast +from areal import workflow_context from areal.api.cli_args import ( GenerationHyperparameters, GRPOConfig, @@ -17,7 +18,6 @@ from areal.api.io_struct import ModelRequest, ModelResponse from areal.api.reward_api import AsyncRewardWrapper from areal.api.workflow_api import RolloutWorkflow -from areal.core import workflow_context from areal.utils import logging, stats_tracker from prompts import ANSWER, SYSTEM_PROMPT, TORL_PROMPT # isort: skip diff --git a/recipe/AEnt/gsm8k_aent_grpo.py b/recipe/AEnt/gsm8k_aent_grpo.py index 2a1b630652..acbf9a61ec 100644 --- a/recipe/AEnt/gsm8k_aent_grpo.py +++ b/recipe/AEnt/gsm8k_aent_grpo.py @@ -8,13 +8,13 @@ from recipe.AEnt.actor import FSDPAEntPPOActor from recipe.AEnt.aent_args import AEntGRPOConfig +from areal import current_platform from areal.api.alloc_mode import AllocationMode from areal.api.cli_args import load_expr_config from areal.api.io_struct import FinetuneSpec, StepInfo, WeightUpdateMeta from areal.dataset import get_custom_dataset from areal.engine.fsdp_engine import FSDPPPOActor from areal.engine.sglang_remote import RemoteSGLangEngine -from areal.platforms import current_platform from areal.reward.gsm8k import gsm8k_reward_fn from areal.utils import seeding, stats_tracker from areal.utils.data import (