feat(ppo): support actor loss aggregation modes - #1
Conversation
4b57f8a to
840c2a6
Compare
|
Important Review skippedToo many files! This PR contains 320 files, which is 170 over the limit of 150. To get a review, narrow the scope: Upgrade to Pro+ to raise the limit. Usage-priced reviews support at most 300 files. ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: ⛔ Files ignored due to path filters (4)
📒 Files selected for processing (324)
You can disable this status message by setting the 📝 WalkthroughWalkthroughThe PR adds ChangesLoss reduction and grouping
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
93af9b7 to
076ad74
Compare
|
@coderabbitai review |
✅ Action performedReview finished.
|
076ad74 to
d109cf2
Compare
d109cf2 to
abe0eab
Compare
There was a problem hiding this comment.
Actionable comments posted: 8
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
areal/utils/functional/functional.py (1)
230-233: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winDon't collapse
±inflog-ratios to zero.This also rewrites real infinite ratios into a neutral ratio of
1, so tokens where one policy assigns zero probability can slip through mask/clamp checks instead of being rejected. OnlyNaNfrom cases like-inf - (-inf)should be neutralized.Proposed fix
- log_ratio = torch.where(torch.isfinite(log_ratio), log_ratio, 0.0) + log_ratio = torch.where(torch.isnan(log_ratio), torch.zeros_like(log_ratio), log_ratio)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@areal/utils/functional/functional.py` around lines 230 - 233, The `log_ratio` sanitization in `functional.py` is collapsing both `NaN` and real `±inf` values to zero, which incorrectly neutralizes valid infinite ratios. Update the `log_ratio` post-processing so only non-finite values produced by invalid same-sign infinities like `-inf - (-inf)` are replaced, while preserving true `±inf` results from `proximal_logprobs` and `old_logprobs`. Keep the fix localized to the `log_ratio` computation block and ensure downstream mask/clamp checks still see real infinite ratios.areal/experimental/engine/archon_engine.py (1)
541-554: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftInclude CP ranks in the training normalizer path.
When CP/Ulysses is enabled,
_gather_actor_train_outputs()gathers acrossself._cp_group, so every CP rank computes the gathered loss. Reducingglobal_normalizersover onlyself.data_parallel_groupmisses the CP factor that Megatron explicitly cancels with a DP+CP normalizer group, causing CP training gradients to be scaled incorrectly.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@areal/experimental/engine/archon_engine.py` around lines 541 - 554, The training normalizer path in `process_output` is only using `compute_global_normalizers(..., self.data_parallel_group)`, which misses the CP/Ulysses factor when `_gather_actor_train_outputs()` has gathered across `self._cp_group`. Update the normalizer computation in `ArchonEngine` to include the CP ranks in the group used for `global_normalizers` (matching the DP+CP normalizer behavior), and keep the rest of `_compute_logprobs_and_loss` unchanged.
🧹 Nitpick comments (1)
tests/experimental/training_service/test_worker_unit.py (1)
99-103: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winExercise the serialized
LossReductioncontract here.Passing
loss_reduction=Nonekeeps the route green, but it no longer validates the new train-batch wire contract this PR introduces. Add one worker test that posts a real serializedLossReduction.mean(...)payload end-to-end so RPC serialization/deserialization regressions are caught.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/experimental/training_service/test_worker_unit.py` around lines 99 - 103, The worker test currently uses loss_reduction=None, which bypasses the new train-batch wire contract; update the relevant test in test_worker_unit to send a real serialized LossReduction.mean payload through the same RPC path. Use the existing serialize_value helper and the train-batch request/response flow in the worker test to verify serialization and deserialization end-to-end, so regressions in LossReduction handling are caught.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@areal/engine/core/train_engine.py`:
- Around line 62-65: The global normalizer check in the training engine is too
strict and causes a crash when every rank contributes zero. Update the
loss-normalization logic around dist.all_reduce(normalizers, group=dp_group) so
fully empty global normalizers are allowed instead of asserting positivity; keep
the zero-handling path consistent with _scale_loss_term(), which already masks
global_normalizer == 0. Adjust the assertion or replace it with logic that
returns a zero loss when all normalizers are zero.
In `@areal/experimental/training_service/controller/controller.py`:
- Around line 715-717: The JSON RPC payload in controller logic is currently
serializing a LossReduction object that may contain non-serializable callables
via loss_fn, which can break encoding before the worker receives it. Update the
payload construction in the controller methods that build the request to send
only a serializable loss descriptor or registry key instead of the LossReduction
instance. Then reconstruct LossReduction on the worker side using that
descriptor so the payload remains JSON-safe. Use the existing payload-building
code paths in controller methods around the batch request to locate both
occurrences.
In `@areal/trainer/ppo/actor.py`:
- Around line 362-395: The M2PO normalizer is being computed from the pre-filter
mask, so mean reduction can weight microbatches using the wrong denominator when
grpo_loss_fn later narrows loss_mask via m2_threshold. Update the normalization
path in actor.py around the LossReduction.mean/ sum setup and the grpo_loss_fn
mask filtering so the normalizer is derived from the post-mask token
count/denominator used after filtering, using the existing symbols
compute_local_normalizers(), grpo_loss_fn, and loss_mask to keep aggregation
consistent across microbatches and ranks.
- Around line 345-356: The minibatch splitting in actor.py is still using
group_size as a hard granularity for prompt_mean, which forces exact
divisibility and breaks on partially filtered groups. Update the split in the
actor minibatch preparation path, around split_padded_tensor_dict_into_mb_list
and MicroBatchSpec, to use a partial-group-aware boundary strategy instead of
requiring bs % granularity == 0. Keep the prompt_mean behavior intact while
allowing incomplete valid groups to pass through so the new partial-group
normalizers can handle them later.
In `@areal/utils/functional/functional.py`:
- Line 1210: The CISPO loss path is still hard-coded to token-level aggregation
because the `ciso...`/CISPO branch in `functional.py` only uses `return_sum` and
never accepts or forwards `loss_aggregation`, `group_size`, or
`loss_aggregation_divisor`. Update the CISPO loss function and its call site in
`areal/trainer/ppo/actor.py` to take the same aggregation parameters as
`ppo_actor_loss_fn` and `sapo_loss_fn`, then route the reduction through
`aggregate_pg_loss` / `aggregate_pg_loss_sum` so `seq_mean`, `prompt_mean`, and
`constant` work when `use_cispo_loss=True`.
In `@tests/fp8/model_hooks.py`:
- Around line 157-159: Use the same context-parallel-aware normalizer group as
MegatronEngine.train_batch in the test helper. In the compute_global_normalizers
call inside model_hooks.py, replace the plain mpu.get_data_parallel_group()
usage with the CP-aware variant so the mirrored path matches training on CP runs
and the gathered gradients stay aligned.
In `@tests/test_tree_training.py`:
- Around line 315-324: The mean-loss test is using a local scalar from
logprobs.mean() that is normalized differently than normalizer_fn, which can
create false gradient drift when padding or masked tokens differ. Update the
loss setup in test_tree_training.py so the scalar passed to LossReduction.mean()
is normalized with the same denominator as normalizer_fn, using the same
active-token count derived from loss_mask in loss_fn and normalizer_fn. Keep the
change localized around loss_fn, normalizer_fn, and LossReduction.mean().
In `@tests/torchrun/run_megatron_engine_vlm_distributed.py`:
- Around line 283-288: The smoke test’s LossReduction.mean setup is using a
constant normalizer, which makes uneven microbatches contribute equally instead
of by size. Update the loss reduction in run_megatron_engine_vlm_distributed.py
to use a size-based normalizer in the LossReduction.mean call, keeping the
existing loss_fn but changing the normalizer_fn to derive the batch/microbatch
size from the input data so the test matches the full-batch loss contract.
---
Outside diff comments:
In `@areal/experimental/engine/archon_engine.py`:
- Around line 541-554: The training normalizer path in `process_output` is only
using `compute_global_normalizers(..., self.data_parallel_group)`, which misses
the CP/Ulysses factor when `_gather_actor_train_outputs()` has gathered across
`self._cp_group`. Update the normalizer computation in `ArchonEngine` to include
the CP ranks in the group used for `global_normalizers` (matching the DP+CP
normalizer behavior), and keep the rest of `_compute_logprobs_and_loss`
unchanged.
In `@areal/utils/functional/functional.py`:
- Around line 230-233: The `log_ratio` sanitization in `functional.py` is
collapsing both `NaN` and real `±inf` values to zero, which incorrectly
neutralizes valid infinite ratios. Update the `log_ratio` post-processing so
only non-finite values produced by invalid same-sign infinities like `-inf -
(-inf)` are replaced, while preserving true `±inf` results from
`proximal_logprobs` and `old_logprobs`. Keep the fix localized to the
`log_ratio` computation block and ensure downstream mask/clamp checks still see
real infinite ratios.
---
Nitpick comments:
In `@tests/experimental/training_service/test_worker_unit.py`:
- Around line 99-103: The worker test currently uses loss_reduction=None, which
bypasses the new train-batch wire contract; update the relevant test in
test_worker_unit to send a real serialized LossReduction.mean payload through
the same RPC path. Use the existing serialize_value helper and the train-batch
request/response flow in the worker test to verify serialization and
deserialization end-to-end, so regressions in LossReduction handling are caught.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: ddd5e05d-7015-41d1-af26-65590355bc89
📒 Files selected for processing (41)
areal/api/__init__.pyareal/api/cli_args.pyareal/api/engine_api.pyareal/engine/core/__init__.pyareal/engine/core/train_engine.pyareal/engine/fsdp_engine.pyareal/engine/megatron_engine.pyareal/experimental/engine/archon_engine.pyareal/experimental/training_service/controller/controller.pyareal/infra/remote_inf_engine.pyareal/trainer/dpo/dpo_engine.pyareal/trainer/ppo/actor.pyareal/trainer/ppo/critic.pyareal/trainer/rw/rw_engine.pyareal/trainer/sft/lm_engine.pyareal/utils/data.pyareal/utils/functional/__init__.pyareal/utils/functional/functional.pydocs/en/best_practices/perf_profiling.mddocs/en/cli_reference.mddocs/generate_cli_docs.pydocs/zh/best_practices/perf_profiling.mddocs/zh/cli_reference.mdtests/experimental/archon/torchrun/run_archon_engine_pp.pytests/experimental/training_service/fake_train_engine.pytests/experimental/training_service/test_worker_unit.pytests/fp8/model_hooks.pytests/test_cispo_loss.pytests/test_dpo.pytests/test_eval_dispatch.pytests/test_grouped_rollout_min_valid.pytests/test_loss_reduction.pytests/test_megatron_engine.pytests/test_partial_group_norm.pytests/test_prompt_mean_loss.pytests/test_train_engine.pytests/test_tree_training.pytests/torchrun/run_fsdp_dcp_distributed.pytests/torchrun/run_fsdp_ulysses_train_batch.pytests/torchrun/run_megatron_engine_distributed.pytests/torchrun/run_megatron_engine_vlm_distributed.py
abe0eab to
99f85bf
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
areal/experimental/engine/archon_engine.py (1)
582-607: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winUse the same group for eval normalization and final aggregation.
eval_batch()normalizes withself.loss_normalizer_group(dp_cp) but still aggregates the final scalar overself.data_parallel_grouponly. Since_gather_actor_train_outputs()reconstructs full outputs on every CP rank, CP-enabled eval will report a loss smaller bycontext_parallel_size.Possible fix
return aggregate_eval_losses( losses if self.pp_has_last_stage else None, - self.data_parallel_group, + self.loss_normalizer_group, self.pp_has_last_stage, self.parallel_dims.get_group("pp") if self.parallel_dims.pp_enabled else None, self._pp_last_stage_rank, )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@areal/experimental/engine/archon_engine.py` around lines 582 - 607, `eval_batch()` is normalizing losses with `self.loss_normalizer_group` but still reducing the final result with `aggregate_eval_losses()` over `self.data_parallel_group` only, which breaks CP-enabled eval. Update the aggregation path in `eval_batch()` to use the same normalization group as `compute_global_normalizers()` and `_gather_actor_train_outputs()` (the `dp_cp` group), keeping the `forward_backward_batch()` and `aggregate_eval_losses()` calls consistent with `self.loss_normalizer_group` and the existing `parallel_dims`/`pp` handling.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@areal/infra/rpc/serialization.py`:
- Around line 272-291: The SerializedCallable.from_callable/to_callable path is
unsafe because it round-trips arbitrary callables through ray.cloudpickle, which
can lead to worker-side code execution when serialize_value() and
deserialize_value() handle RPC payloads. Replace the cloudpickle-based
serialization with an allowlisted mechanism in SerializedCallable, such as
storing a vetted import path or registry key and resolving it only against
known-safe callables in to_callable(). Ensure the serialized data in
SerializedCallable and the caller paths in serialize_value()/deserialize_value()
only accept explicitly approved callable identifiers.
In `@areal/utils/functional/functional.py`:
- Around line 1024-1025: The PG loss reduction logic currently converts
loss_mask and denom_mask to boolean masks without verifying their shapes against
pg_loss, so add explicit shape checks before any masking/reduction in the helper
that computes num_mask and den_mask. In the PG loss path(s) where these masks
are applied, reject mismatched shapes rather than relying on broadcasting, and
ensure both loss_mask and denom_mask match pg_loss before they are used to
compute numerator and denominator values.
In `@tests/test_prompt_mean_loss.py`:
- Line 576: The pytest assertion in the prompt mean loss test is using a regex
match string that treats gconfig.n_samples as a pattern, so the dot is currently
unescaped. Update the ValueError expectation in the relevant pytest.raises call
to escape the dot in the match text so it matches the literal message, and keep
the change localized to the test around the prompt mean loss validation.
---
Outside diff comments:
In `@areal/experimental/engine/archon_engine.py`:
- Around line 582-607: `eval_batch()` is normalizing losses with
`self.loss_normalizer_group` but still reducing the final result with
`aggregate_eval_losses()` over `self.data_parallel_group` only, which breaks
CP-enabled eval. Update the aggregation path in `eval_batch()` to use the same
normalization group as `compute_global_normalizers()` and
`_gather_actor_train_outputs()` (the `dp_cp` group), keeping the
`forward_backward_batch()` and `aggregate_eval_losses()` calls consistent with
`self.loss_normalizer_group` and the existing `parallel_dims`/`pp` handling.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 0e1da45c-49a8-46b0-a89d-e44f9aa74582
📒 Files selected for processing (42)
areal/api/__init__.pyareal/api/cli_args.pyareal/api/engine_api.pyareal/engine/core/__init__.pyareal/engine/core/train_engine.pyareal/engine/fsdp_engine.pyareal/engine/megatron_engine.pyareal/experimental/engine/archon_engine.pyareal/experimental/training_service/controller/controller.pyareal/infra/remote_inf_engine.pyareal/infra/rpc/serialization.pyareal/trainer/dpo/dpo_engine.pyareal/trainer/ppo/actor.pyareal/trainer/ppo/critic.pyareal/trainer/rw/rw_engine.pyareal/trainer/sft/lm_engine.pyareal/utils/data.pyareal/utils/functional/__init__.pyareal/utils/functional/functional.pydocs/en/best_practices/perf_profiling.mddocs/en/cli_reference.mddocs/generate_cli_docs.pydocs/zh/best_practices/perf_profiling.mddocs/zh/cli_reference.mdtests/experimental/archon/torchrun/run_archon_engine_pp.pytests/experimental/training_service/fake_train_engine.pytests/experimental/training_service/test_worker_unit.pytests/fp8/model_hooks.pytests/test_cispo_loss.pytests/test_dpo.pytests/test_eval_dispatch.pytests/test_grouped_rollout_min_valid.pytests/test_loss_reduction.pytests/test_megatron_engine.pytests/test_partial_group_norm.pytests/test_prompt_mean_loss.pytests/test_train_engine.pytests/test_tree_training.pytests/torchrun/run_fsdp_dcp_distributed.pytests/torchrun/run_fsdp_ulysses_train_batch.pytests/torchrun/run_megatron_engine_distributed.pytests/torchrun/run_megatron_engine_vlm_distributed.py
✅ Files skipped from review due to trivial changes (5)
- areal/engine/core/init.py
- docs/en/cli_reference.md
- docs/en/best_practices/perf_profiling.md
- docs/zh/best_practices/perf_profiling.md
- docs/zh/cli_reference.md
🚧 Files skipped from review as they are similar to previous changes (26)
- areal/api/init.py
- tests/test_megatron_engine.py
- areal/trainer/ppo/critic.py
- areal/utils/functional/init.py
- tests/torchrun/run_fsdp_dcp_distributed.py
- areal/trainer/sft/lm_engine.py
- tests/test_loss_reduction.py
- areal/trainer/dpo/dpo_engine.py
- tests/test_train_engine.py
- tests/test_tree_training.py
- tests/test_partial_group_norm.py
- docs/generate_cli_docs.py
- tests/experimental/archon/torchrun/run_archon_engine_pp.py
- areal/infra/remote_inf_engine.py
- tests/test_dpo.py
- tests/torchrun/run_megatron_engine_distributed.py
- areal/api/engine_api.py
- areal/trainer/rw/rw_engine.py
- tests/test_grouped_rollout_min_valid.py
- areal/experimental/training_service/controller/controller.py
- areal/engine/core/train_engine.py
- tests/torchrun/run_fsdp_ulysses_train_batch.py
- areal/engine/megatron_engine.py
- areal/engine/fsdp_engine.py
- tests/fp8/model_hooks.py
- areal/trainer/ppo/actor.py
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@areal/infra/rpc/serialization.py`:
- Around line 49-67: The RPC callable registry is only populated in tests, so
production receivers cannot deserialize callable payloads. Update the
worker/bootstrap startup path to import the module(s) that call
register_rpc_callable(), or otherwise register the needed functions before any
serialize_value()/deserialize_value() RPC traffic begins. Use the existing
_RPC_CALLABLE_REGISTRY, register_rpc_callable(), and
_registered_rpc_callable_key() flow to ensure the same callable keys are
available in both sender and receiver processes.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 02f33a70-11b5-4654-b917-eff9e3939a5b
📒 Files selected for processing (4)
areal/infra/rpc/serialization.pyareal/utils/functional/functional.pytests/test_prompt_mean_loss.pytests/test_rpc_callable_serialization.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/test_prompt_mean_loss.py
- areal/utils/functional/functional.py
5362665 to
f95c73b
Compare
areal-project#1440) * feat(cli): add experimental cli scaffold for service-style subcommands Prepares the ground for feat/inference-service-cli, feat/training-service-cli, and feat/agent-service-cli to land on top — each subcommand PR only adds its own click group and one cli.add_command call, instead of duplicating the same plumbing. Key pieces: - empty top-level `cli` click group with version + --help wiring; subcommand modules attach themselves via cli.add_command(...) - `state.areal_home()` + `atomic_write_json()` for AREAL_HOME-rooted local state files (atomic .tmp + os.replace) - `process` module: pid_alive / pick_free_port / signal_pid / kill_pids + spawn_process with start_new_session=True so detached worker processes survive parent SIGHUP and a graceful SIGTERM→grace→SIGKILL teardown sequence is one helper away - pyproject: `[project.scripts] areal = ...cli.main:cli`, optional `[cli]` deps (click, colorlog) Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * build(cli): regenerate uv lockfiles for cli optional deps Pre-commit's uv-lock hook expects the lockfiles to track click and colorlog after they were added to [project.optional-dependencies] cli. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * fix(cli): address gemini review on process/state utilities - pid_alive: try os.waitpid(pid, WNOHANG) before os.kill so a zombie child does not look alive — otherwise kill_pids waits its full grace window and sends a redundant SIGKILL. - spawn_process: close the parent's log_handle copy after Popen dup()s it for the child, so repeated spawns do not leak fds. - atomic_write_json: use tempfile.NamedTemporaryFile + os.fsync, and unlink the tempfile on serialization or rename failure — fixes a tempfile leak on bad input, durability against crashes, and a race between concurrent writers on the same path. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * refactor(cli): unify terminate + reuse kill_process_tree, drop pick_free_port Replace the scaffold's bespoke process primitives with shared utilities already in the codebase: - process: kill_pids now delegates to areal.infra.utils.proc.kill_process_tree, which walks the descendant tree via psutil. This fixes the bug TaoZex raised on signal_pid's inconsistent PermissionError handling — signal_pid is gone entirely. - process: drop pick_free_port. Callers should use areal.utils.network.find_free_ports, which draws from a non-ephemeral port range and supports exclude_ports — addresses the TOCTOU race guozhihao-224 flagged on the naive bind(0) implementation. - scheduler: new module with a stateless top-level terminate(ref, *, backend, grace_s) that subcommand CLIs can funnel teardown through. Only the local backend is implemented; future backends extend the dispatch in-module. The free-function shape keeps the dispatcher from drifting into stateful scheduler instances. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * chore: integrate CLI shared files Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * style(cli): drop extra blank line after imports in utils.py Local pre-commit auto-fixed the blank line but the fix wasn't re-staged before commit 9401918, so CI's ruff hook caught it. * refactor(cli): collect namespace path helpers into NamespacedStateStore state.py previously exposed 13 free functions (namespace_root, services_dir, logs_dir, service_state_path, …, recover_pids_from_raw_state) that all took ``namespace`` as their first argument. Promote them to methods on a new NamespacedStateStore class so subcommand CLIs construct one instance per namespace and stop threading the namespace string through every call. areal_home, atomic_write_json, SupportsComponentProbe, and ServiceStateBase remain at module level. ServiceLifecycle and ConfigLoader each hold a store instance built from their namespace at __init__ time; LogsCommand resolves log paths via ``lifecycle.store.logs_dir``. No behavior change. Note: downstream branches (inf / agent / train) will need to swap their imports from ``state.service_state_path``/``state.clear_current_service``/ etc. to ``store.<method>`` on the next rebase. --------- Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
…real-project#1448) * refactor: move 5 experimental modules into areal/v2 for 2.0 release Move agent_service, inference_service, training_service, weight_update, and cli from areal/experimental/ to areal/v2/, and rewrite every reference (Python imports, `python -m` invocations, console-script entry points, CODEOWNERS, docs, review-pr signals) to the new path. - 5 directories migrated via `git mv` (history preserved) - 70 files modified across areal/, tests/, examples/, pyproject{,vllm}.toml, .github/CODEOWNERS, ROADMAP.md, and tooling docs - `areal/v2/__init__.py` added so v2 is an importable package - Build config (tool.uv.build-backend with module-root="") auto-discovers the new package — no pyproject changes beyond the entry point Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * style: apply pre-commit auto-fixes after areal/v2 move CI pre-commit job reformatted 54 files automatically: - .github/CODEOWNERS: realign owner columns (32-col) after shorter /areal/v2/ paths broke the previous 40-col alignment - areal/**, tests/**, examples/**, docs/**: ruff isort reorders `areal.v2.*` imports into their new alphabetical slot (between `areal.engine` and `areal.infra`) - markdown/yaml whitespace normalized by mdformat / ruff-format Pure formatting; no logic change. `pre-commit run --all-files` is now green locally. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * fix(tests): point py<3.12 conftest stubs at areal/v2/ Two test conftest.py stubs (Python 3.10/3.11 compat) still had the `areal/experimental/{inference_service,weight_update}` path as the namespace package's __path__, broken by the v2 move. The sed pass missed them because the paths were comma-separated `os.path.join` args (`"areal", "experimental", "X"`), not slash-form path strings. Additionally insert an `areal.v2` stub between `areal` and the leaf package so the parent→child attribute wiring loop (which uses `name.rsplit(".", 1)`) can find a parent module in sys.modules. Without it `setattr(parent, child, ...)` silently no-ops and `unittest.mock.patch` traversal breaks on the new path. Spotted by gemini-code-assist on PR areal-project#1448. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * refactor(tests): mirror areal/v2 move under tests/v2 Move agent_service, inference_service, training_service, and weight_update test directories from tests/experimental/ to tests/v2/ to mirror the source-tree layout. tests/experimental/ retains archon/ and openai/ (still-experimental modules). - 50 files migrated via `git mv` (history preserved) - tests/v2/__init__.py added - 9 files rewritten for the new dotted/slashed test paths: `tests.experimental.{4 modules}` → `tests.v2.{4 modules}` `tests/experimental/{4 modules}` → `tests/v2/{4 modules}` (covers integration_utils imports, fake_train_engine engine_class, pytest invocation strings in docstrings, and the tests/v2/weight_update/torchrun/run_nccl_weight_transfer.py path) conftest stubs (Python <3.12 namespace shims) keep working because _REPO_ROOT is computed via "..", "..", ".." from the conftest file — same depth under tests/v2/X/ as under tests/experimental/X/. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
* chore: integrate CLI shared files
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* feat(cli): training experiment launcher — `areal train run`
Wraps a user-supplied driver entry point (`module.path:func`) into a
uniform CLI invocation driven by yaml + hydra overrides. Same lifecycle
as `python -m examples.math.gsm8k_rl ...` — the CLI process is the
driver process, attached, all the way down.
Single verb, permanently. No ps / stop / status / logs: a training job
is a process, not a service; OS tooling (`ps`, `kill`) and the cluster
scheduler (`squeue`, `kubectl get`) already cover that, and logs live
at `{fileroot}/logs/...` for `tail -F`. Service-shaped verbs belong
with `areal inf`, whose state is daemon-backed and invisible to OS
tooling.
Usage:
areal train run --config experiments/grpo.yaml \
--driver examples.math.gsm8k_rl:main \
actor.lr=1e-5 trial_name=lr-sweep-3 +actor._version=v2
`--driver MOD:FN` is required; we deliberately don't peek a `driver:`
field out of the yaml — keeps the contract obvious and avoids parsing
yaml in the CLI layer.
Anything after the flags is forwarded verbatim as a hydra override to
the driver (click's `nargs=-1 + UNPROCESSED + ignore_unknown_options`).
Backgrounding is out of scope: users wrap the command in
nohup/tmux/sbatch as appropriate. Adding `--detach` would force state
files, heartbeat, pid tracking — none of which anything else here needs.
Conflicts with `feat/inference-service-cli` are isolated to main.py +
pyproject.toml's [cli] extra; trivial 3-way merge after the inf branch
lands.
* refactor(experimental): split training CLI commands
Move the training CLI entrypoint into a root CLI module and a training command package so future training commands can be added without growing a monolithic command file.
* docs(cli/train): add training service CLI guide
Document the `areal train` subcommand group: basic concepts, the
driver function contract, hydra override conventions, and exit code
behaviour.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* chore(experimental): rebase train CLI onto cli scaffold v2
Train doesn't carry a service / daemon / state, so the scaffold's
ServiceLifecycle / BaseHTTPClient / StatusReporter components don't
apply — train run is a synchronous importlib dispatcher. The change
here is just plumbing alignment:
- Rebase onto feat/experimental-cli-scaffold; scaffold now owns
cli/__main__.py / cli/cli.py.
- Drop the redundant cli/main.py shim (cli/__main__.py covers it).
- Point [project.scripts] areal entry at cli.cli:cli in both
pyproject manifests, matching the agent / inf branches.
- Re-flow training/cli_guide.md through mdformat.
* docs(cli/train): drop the evaluation example from the guide
* docs(cli/train): drop the BOBA-GRPO example from the guide
* fix(cli): resolve config merge leftovers
---------
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
f95c73b to
b2a18f0
Compare
* feat(cli): add experimental cli scaffold for service-style subcommands Prepares the ground for feat/inference-service-cli, feat/training-service-cli, and feat/agent-service-cli to land on top — each subcommand PR only adds its own click group and one cli.add_command call, instead of duplicating the same plumbing. Key pieces: - empty top-level `cli` click group with version + --help wiring; subcommand modules attach themselves via cli.add_command(...) - `state.areal_home()` + `atomic_write_json()` for AREAL_HOME-rooted local state files (atomic .tmp + os.replace) - `process` module: pid_alive / pick_free_port / signal_pid / kill_pids + spawn_process with start_new_session=True so detached worker processes survive parent SIGHUP and a graceful SIGTERM→grace→SIGKILL teardown sequence is one helper away - pyproject: `[project.scripts] areal = ...cli.main:cli`, optional `[cli]` deps (click, colorlog) Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * refactor(cli): unify terminate + reuse kill_process_tree, drop pick_free_port Replace the scaffold's bespoke process primitives with shared utilities already in the codebase: - process: kill_pids now delegates to areal.infra.utils.proc.kill_process_tree, which walks the descendant tree via psutil. This fixes the bug TaoZex raised on signal_pid's inconsistent PermissionError handling — signal_pid is gone entirely. - process: drop pick_free_port. Callers should use areal.utils.network.find_free_ports, which draws from a non-ephemeral port range and supports exclude_ports — addresses the TOCTOU race guozhihao-224 flagged on the naive bind(0) implementation. - scheduler: new module with a stateless top-level terminate(ref, *, backend, grace_s) that subcommand CLIs can funnel teardown through. Only the local backend is implemented; future backends extend the dispatch in-module. The free-function shape keeps the dispatcher from drifting into stateful scheduler instances. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * chore: integrate CLI shared files Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * refactor(cli): collect namespace path helpers into NamespacedStateStore state.py previously exposed 13 free functions (namespace_root, services_dir, logs_dir, service_state_path, …, recover_pids_from_raw_state) that all took ``namespace`` as their first argument. Promote them to methods on a new NamespacedStateStore class so subcommand CLIs construct one instance per namespace and stop threading the namespace string through every call. areal_home, atomic_write_json, SupportsComponentProbe, and ServiceStateBase remain at module level. ServiceLifecycle and ConfigLoader each hold a store instance built from their namespace at __init__ time; LogsCommand resolves log paths via ``lifecycle.store.logs_dir``. No behavior change. Note: downstream branches (inf / agent / train) will need to swap their imports from ``state.service_state_path``/``state.clear_current_service``/ etc. to ``store.<method>`` on the next rebase. * feat(experimental): add agent service CLI * refactor(experimental): align agent CLI with Click * refactor(experimental): remove agent rl negotiated flag Use the existing RL session id and API key as the source of truth instead of keeping a derived boolean in agent session state. Key changes: - Drop rl_negotiated from SessionState and tolerate legacy state files - Route agent CLI and demo output through loggers instead of print - Update CLI state tests for the revised session schema * refactor(experimental): avoid creating agent session on run Agent service startup should only launch runtime components. Sessions are now created explicitly through new_session. Key changes: - Remove initial session creation from areal agent run - Drop the run --session-key option - Add CLI tests for session-free startup * refactor(experimental): align agent CLI with inference; drop session verbs Pulled the agent CLI in line with feat/inference-service-cli and removed surface the controller-replacement scope does not need. Style alignment with inference CLI: - adopt 2-step *_cmd -> do_* (drop the redundant handle layer) - table / JSON output via click.echo; hard failures via click.ClickException - Click default_map driven by load_click_default_map(_BINDINGS); drop the hand-rolled cfg_get / resolve_* helpers - single AgentCli logger (registered color in areal/utils/logging.py), exported from new agent/common.py alongside running_state / load_running_state / wait_http_health - AgentCLIHTTPError / AgentCLIUnreachable renamed to AgentHTTPError / AgentUnreachable; per-server clients renamed (GatewayClient, RouterClient, DataProxyClient) to mirror inference/client.py Removed surface: - interactive REPL stub (interactive.py + --interactive / --stop-on-exit / --history-file on run); chat / reward placeholders never shipped - destroy / health aliases of stop / status - new_session / switch_session and session_ops; SessionState / SessionsState / InferenceClient go with them. Session lifecycle is moving to a gateway /v1/sessions REST endpoint; CLI no longer owns it - legacy rl_negotiated migration in state.py (the flag was removed two commits ago; no production state needs the shim) Net diff: -830 lines, 6 tests green, pre-commit clean. * refactor(experimental): drop cli scaffold overlap pending areal-project#1440 The shared cli/cli.py + cli/main.py entry points and the pyproject [project.scripts] + the example invocation tweak duplicated plumbing that the cli scaffold PR (areal-project#1440) now provides. Drop them here so this branch can rebase cleanly on top of areal-project#1440 once that lands; the agent CLI itself (cli/agent/*) is untouched and continues to live in this PR. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * refactor(experimental): rename agent http.py to client.py for parity with inf inf-cli puts the same kind of HTTP client wrappers in cli/inference/client.py; rename the agent counterpart to match so the two sibling subcommand packages have parallel layouts. Pure rename + three import updates, no behavior change. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * refactor(experimental): adopt cli scaffold v2 in agent service Rebase onto feat/experimental-cli-scaffold and replace the per-CLI duplicates with the shared base classes/utilities: - state.py now keeps only the agent dataclasses; ServiceState satisfies ServiceStateBase and gateway_alive() is anchored on the gateway PID only (fixes the prior "any pid alive" check that kept reporting running after the gateway died). - config.py / client.py / lifecycle.py shrink to thin subclasses of ConfigLoader / BaseHTTPClient / ServiceLifecycle. - run.py and stop.py route through ServiceLifecycle for double-start refusal, force-replace, and state cleanup. - status.py renders via StatusReporter + ColumnSpec; ps.py emits via json_or_table; the bespoke logs subcommand is dropped in favor of LogsCommand(lifecycle=...).build(). - launcher.py reserves all ports up front through find_free_ports (non-ephemeral, no TOCTOU) and uses the scaffold spawn/wait helpers. Net ~360 LOC dropped from the agent module while every behavioral guarantee carries over. * refactor(experimental): drop empty Gateway/DataProxy client subclasses Both classes only inherited BaseHTTPClient.health() and were never instantiated after the scaffold refactor (StatusReporter probes addrs directly, launcher only uses RouterClient). Callers that ever need a gateway/data-proxy health probe can do BaseHTTPClient(addr) — the variable name carries the role. RouterClient stays because register_proxy is a real method. * fix(experimental): point areal CLI script at cli.py entry main.py was dropped earlier when scaffold landed (its single-line shim was redundant with __main__.py), but the [project.scripts] entry still referenced it, leaving the installed ``areal`` console script broken. Update both pyproject manifests to ``cli.cli:cli``. * docs(cli/agent): add agent service CLI guide Document the `areal agent` subcommand group: launching the gateway/router with N worker/data-proxy pairs, inspecting service state, log management, two-phase shutdown, configuration precedence, and the relationship with `areal inf`. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * style(docs): apply mdformat to agent CLI guide * refactor(experimental): adopt scaffold NamespacedStateStore in agent scaffold's state.py promoted the namespace-aware free functions onto NamespacedStateStore. Update agent accordingly: hold a module-level ``store = NamespacedStateStore(AGENT_NAMESPACE)`` and route every service_state_path / set_current_service / clear_current_service / logs_dir call through it. No subclassing needed — agent has no extra state files (vs inf's two-file split). Behavior unchanged; 6 tests still pass. * style: sort v2 CLI imports * refactor(v2/cli/agent): drop dead --inf-* options The --inf-addr / --inf-api-key / --inf-model surface was a placeholder for a session-api-key negotiation that never landed: the values were written into ServiceState but no downstream code (launcher, gateway, worker, data_proxy, /v1/responses bridge) reads them. Keeping the options advertises functionality the CLI does not provide. - run.py: drop the three click options and do_run params - launcher.py: drop the inf_* kwargs on launch_agent_stack - state.py: drop inf_* fields from ServiceState; pop them on load() for forward-compat with state files written by older revisions - config.py: drop the [inference] -> run.inf_* TOML bindings - cli_guide.md: drop the --inf-* example, the "Relationship with areal inf" section, and the [inference] TOML block; also drop the "Multiple pairs" example and fix the agent import path syntax from the package:Class colon form to the module.Class dot form (the worker resolves the path via import_from_string which requires dots) - tests: drop inf_* kwargs from test_agent_cli_run; rewrite the config default_map test to assert on a non-inference key Will be reintroduced (with real wiring) when /v1/sessions lands on the gateway and the gateway process needs inf credentials to negotiate session_api_key for RL trajectory association. --------- Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
* feat(cli): inference service CLI — daemon + 4 verbs
`areal inf` is an ollama-style operator console for the local inference
service. One daemon per user, one OpenAI-compatible gateway endpoint,
models registered against it.
Verbs:
inf run start daemon (gateway + router); inline --model registers
an external (--api-url) or internal (--backend / --model-path)
model in the same call
inf ps list registered models
inf status daemon health + model count
inf stop SIGTERM gateway+router, grace period, then SIGKILL
State is a single ~/.areal/inf/state.json (pid + url + admin key +
started_at). Everything else lives in the gateway/router process
memory; CLI is otherwise stateless.
Layout under areal/experimental/cli/:
main.py / state.py shared scaffold (areal_home, pid_alive,
atomic_write_json), to be reused by
`areal train` in a separate PR
commands/inf/__init__.py all four verbs + register helpers in
a single file (one place to scan
what the user can do)
commands/inf/state.py DaemonState dataclass + paths
commands/inf/launcher.py subprocess spawn helpers, including
base_gpu_id support for sglang dp>1
commands/inf/client.py urllib gateway + router HTTP client
Heavy imports (sglang/vllm/torch via areal.api.cli_args) stay lazy —
inside register helpers only — so `areal inf -h` and `areal -h` parse
the click tree without paying for them.
* fix(cli/inf): persist model worker pids in state.json so stop kills them
`_register_internal` returns the list of sglang+data-proxy pids it
spawned, but `_do_run` discarded the return value. state.json then
only held gateway+router pids, so `areal inf stop` left sglang and
data-proxy processes orphaned.
Add a `worker_pids: list[int]` field to DaemonState, save the list
right after register_internal succeeds, and include those pids in
the stop kill set.
Also fix _do_stop short-circuiting on a dead gateway pid — even when
the daemon front-end is dead, model worker pids in state may still
be alive (the original bug pattern this commit fixes). Always try
to kill every pid we know about, then drop the state file.
* fix(cli/inf): route startup messages through getLogger; align admin key default
Startup-time prints (router/gateway pids, replica spawn lines, daemon
ready, foreground / shutdown notices) were going through click.echo,
so they bypassed the AReaL log formatter — no timestamps, no log
level, no PascalCase tag. Switch all of these to
`getLogger("InfCli").info(...)` so output looks like the rest of the
project:
20260611-09:14:02.157 InfCli INFO: starting inference daemon ...
20260611-09:14:02.342 InfCli INFO: router pid=304116 http://...
20260611-09:14:02.910 InfCli INFO: gateway pid=304118 http://...
Register InfCli in LOGGER_COLORS_EXACT under the launcher (blue)
group, since `areal inf run` is conceptually a process launcher.
Reads kept on click.echo on purpose: `inf ps` / `inf status` plain-text
output is structured (tabular or single-field-per-line) and meant to
be consumed by jq / awk in scripts; piping that through the colored
log formatter would break it.
Also flip `--admin-api-key` default from "areal-admin-key" to
"admin-api-key" — matches the existing inference_service convention.
* feat(cli/inf): standalone `register` and `deregister` verbs
Phase 2 — model lifecycle separated from daemon lifecycle.
areal inf register <name> [external | internal flags]
Register a new model against a running daemon. Same flags as
`inf run --model ...`, just attached to an existing service.
areal inf deregister <name> [--grace] [--force]
Drop the model from the router, unregister its proxy workers,
and SIGTERM/SIGKILL the spawned sglang+data-proxy procs.
DaemonState gains a `models: dict[str, ModelEntry]` mapping each
registered model to its (pids, proxy_addrs). This replaces the
previous flat `worker_pids: list[int]` so deregister can target one
model's processes without touching the others. `inf stop` flattens
across all entries and kills the whole set; the foreground / failure
cleanup paths use the same flatten.
`_register_internal` now returns `(pids, proxy_addrs)` rather than
just pids — both are needed at deregister time (router unregister
takes the proxy addr; SIGTERM takes the pid).
External models also get a state entry now (empty pids/addrs) so
`inf deregister` can find them and drop them from the router.
phase 1 verbs unchanged.
* feat(cli/inf): phase 3 — reward + collect verbs
Two verbs that round out the RL data path:
areal inf reward <session_api_key> <reward> [--model X]
thin wrapper around POST /rl/set_reward. set_reward is the
only thing that flips an active conversation into a ready
trajectory, so any data-collection flow needs to call it (even
with a dummy reward=0 just to "flush"). CLI verb is for users
whose agent is in another language / shell / human raters; agent
authors writing python are free to POST directly.
areal inf collect <model> --batch-size N \
[--sessions-out FILE] [--output FILE] \
[--timeout T] [--poll-interval S] \
[--discount D] [--style individual|concat]
client-side batch orchestrator. start_session(group_size=N) ->
hand sessions to the agent (via --sessions-out FILE) -> poll
/export_trajectories every poll-interval seconds, accumulating
unique trajectories until N are collected or timeout fires ->
one final export with remove_session=True for cleanup -> dump
JSONL.
Mirrors the controller's rollout_batch path but moves the wait
loop to the client so gateway / router stay stateless. Agent
lifecycle is intentionally NOT inside collect (no --agent-cmd):
agent runs in its own process and just reads sessions_out.
JSONL is the only output format on purpose -- the gateway already
serializes trajectories to JSON at HTTP boundary, so .pt would mean
re-decoding tensors only to re-encode them; trainers consuming the
output can torch.tensor(x) when they need it.
8 verbs total (run / ps / status / stop / register / deregister /
reward / collect).
* feat(cli/inf): enrich ps/status with model kind / backend / addrs
Phase 2.5 — bring `inf ps` and `inf status` to the design_inf.md
fidelity (sections 11.5 / 11.6) without bringing back the multi-service
concept.
ModelEntry gains four fields:
kind: 'internal' | 'external'
backend: spec string ('sglang:tp=2,dp=2') for internal, '' for external
api_url: external upstream URL, '' for internal
inference_server_addrs: per-replica sglang/vllm URLs (internal)
`_register_internal` returns (pids, proxy_addrs, inf_addrs); _do_run
and _do_register both fill in the new ModelEntry fields.
inf ps now reads state.models (CLI-side truth) instead of polling
gateway /v1/models. Output:
NAME KIND BACKEND WORKERS
qwen3 internal sglang:tp=2,dp=2 2
gpt-4o external - -
inf status switches to a multi-row table per design 11.5:
COMPONENT STATUS ADDR DETAILS
gateway ok http://127.0.0.1:8080 models=2
router ok http://127.0.0.1:..
qwen3 registered internal backend=sglang:tp=2 workers=2
gpt-4o registered external api_url=https://...
JSON output of both verbs follows suit.
Backwards-compat: ModelEntry's new fields all have defaults, so an
old state.json still loads cleanly.
* feat(cli/inf): add `inf models` verb
`inf ps` currently lists registered models (table: NAME / KIND /
BACKEND / WORKERS). Add `inf models` as a more explicit alias —
docker-style "different verb for different resource". `ps` keeps
its current behavior; both call the same _print_models helper so
output is identical.
This isn't part of the multi-service rollback (which we decided not
to do). It's a small ergonomic addition for the single-daemon shape.
* feat(cli/inf): add `logs` verb + TOML config support + help text for proxy/engine args
Three related additions:
1. `inf logs --component NAME [-f] [-n LINES]`
Tail a log under ~/.areal/inf/logs/. Component defaults to 'gateway';
can be 'router' or a full model log basename like 'qwen3-inf-0'.
-f follows (tail -F), -n sets initial line count (default 200).
Exec's tail directly for stream fidelity.
2. TOML config support (design 12)
- Group-level option `areal inf --config FILE` merges FILE on top
of ~/.areal/inf/config.toml (both optional).
- config.py loads TOML via tomllib (py3.11+), maps [default] /
[launch] / [register.internal] / [collect] sections to CLI
option defaults via click's default_map mechanism.
- Precedence: CLI flag > --config > ~/.areal/inf/config.toml >
hardcoded default.
3. Detailed help for --engine-args and --proxy-args
Users couldn't guess what to pass. Now both flags show inline
hints (common sglang knobs / data-proxy flags + defaults). Help
is a single paragraph so click's wrap_text handles terminal width.
10 verbs total (run / stop / ps / status / models / register /
deregister / reward / collect / logs).
* fix(cli/inf): align `collect` flags with design 11.9
Rename + add + drop options on `inf collect` so the surface matches
the design spec exactly:
rename --discount -> --turn-discount
rename --style -> --export-style
add --format json|jsonl (default jsonl)
add --json progress-events flag (placeholder; not implemented)
drop --task-id (always 'cli-collect' internally)
drop --sessions-out (agents query gateway directly)
drop --poll-interval (always 2.0s internally)
JSON output (--format json) emits {tid: interaction, ...} pretty-printed.
JSONL output (default) emits one trajectory per line, each prefixed
with trajectory_id.
config.toml [collect] keys renamed to match.
* fix(cli/inf): three bugs (admin key, gpu collision, sglang request log)
P1. --admin-api-key default reverts from "admin-api-key" to
"areal-admin-key" — matches the v2 inference_service convention
used everywhere else in the codebase.
P2. Registering a second internal model collided with the first on
GPUs 0..tp-1. Cause: base_gpu_id was always r * tp, computed only
against the current model's dp index, ignoring GPUs already used
by previously-registered models. Fix: track a monotonic
`next_gpu_id` cursor on DaemonState and pass it into
`_register_internal` as `base_gpu_id`. Each ModelEntry now
records its (base_gpu_id, gpu_count) so deregister can roll back
the cursor when removing the *last* registered model (preserves
contiguous allocation; doesn't try to coalesce holes in the
middle, which is fine for the v1 use case).
P3. sglang server logs only had model-load output, no chat requests.
Cause: SGLangConfig.log_requests defaults to False. Fix: spawn
sglang with log_requests=True so /chat/completions traffic shows
up under ~/.areal/inf/logs/<model>-inf-N.log. (data-proxy
access logs are off via uvicorn config inside the data-proxy
package itself; that's not under inf CLI's control.)
* refactor(experimental): split inference CLI commands
Move the inference service CLI out of the monolithic commands package and into command-specific modules. Align collect with the session/export flow by returning session keys, polling exports without revoking the router group, and cleaning up at the end.
* feat(experimental): support multiple inference services
* fix(experimental): harden inference CLI lifecycle
Protect inference service state transitions so register and run do not race or leave orphaned processes behind.
Key changes:
- Track engine and proxy PIDs separately for phased shutdown
- Lock model state during register and startup model setup
- Clean up foreground services on SIGTERM and SIGHUP
- Recover raw PIDs before forced service replacement
- Align inference CLI model and session option names with the design
* feat(cli/inf): add scheduler abstraction for worker placement
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* fix(cli/inf): widen probe timeout, parallelize status, drop default_model
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* refactor(experimental): adopt cli scaffold v2 in inference service
Rebase onto feat/experimental-cli-scaffold and replace the per-CLI
duplicates with the shared base classes/utilities:
- state.py keeps the inf-specific dataclasses and the two-file
recover_pids_from_raw_state; RuntimeState now satisfies
ServiceStateBase (gateway_alive + components + .load classmethod).
- config.py / client.py / lifecycle.py shrink to thin subclasses of
ConfigLoader / BaseHTTPClient / ServiceLifecycle. The old
GatewayHTTPError / GatewayUnreachable names are kept as aliases of
ServiceHTTPError / ServiceUnreachable so subcommands swap mechanically.
- InferenceLifecycle overrides force_replace_slot to walk the inf
raw-state helper (which knows about the secondary model-state file)
and to remove both files on cleanup.
- common.py drops scaffold-replaced helpers (running_state /
load_running_state / refuse_if_running / wait_http_health /
wait_client_health / print_services / print_models /
probe_http_health); keeps backend-spec parsing, model registration,
TaskHandle formatters, and terminate_runtime_state (data-flow order
is inf-specific).
- commands/run.py uses ServiceLifecycle for refuse / force-replace and
ForegroundWatcher for the SIGINT/SIGTERM/SIGHUP handling.
- commands/stop / status / ps / models / register / deregister /
reward / collect route through inf_lifecycle; status.py emits via
StatusReporter + ColumnSpec; ps/models via json_or_table.
- launcher.py and scheduler/local.py replace pick_free_port with
find_free_ports (non-ephemeral, no TOCTOU); LocalScheduler tracks
allocated ports across submits.
- commands/logs.py is removed; LogsCommand(lifecycle=inf_lifecycle)
wires the verb in __init__.py.
Net ~290 LOC dropped while every behavioral guarantee carries over.
* docs(cli/inf): add inference service CLI guide
Document the `areal inf` subcommand group: launching the gateway/router,
registering models, RL session flow with rewards, trajectory collection,
log management, and configuration file precedence.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* style(docs): apply mdformat to inference CLI guide
* refactor(experimental): address PR review on inference CLI
Address garrett4wade's review on areal-project#1435:
1+2) Drop the unfinished `collect` verb and the gateway RPCs it
depended on. `commands/collect.py` is removed entirely; the gateway
client no longer exposes `start_session` / `export_trajectories`
(only `set_reward` remains, used by the surviving `reward` verb).
The cli group, config bindings, cli_guide section, and parser tests
are cleaned up accordingly.
3) Replace the bespoke `sglang:tp=2,dp=2` mini-DSL in
`parse_backend_spec` with `ModelAllocation.from_str`, so the CLI
accepts the same grammar as `InferenceEngineConfig.backend` in YAML
configs (`sglang:d4`, `vllm:d2t4`). Help text, doc examples, and
test fixtures are updated to the new form.
4) Collect the free functions in `state.py` (`models_dir`,
`models_state_path`, `models_lock_path`, `locked_model_state`,
`recover_pids_from_raw_state`) into an `InferenceStateStore` class
and route every caller through a module-level `store` instance.
The dataclasses now obtain paths via `store.<...>`, keeping
on-disk-layout responsibilities in one place.
Net: 13 files, +170 / -524 (mostly from dropping the collect verb).
* style(tests): drop trailing blank lines after removed collect case
* refactor(experimental): subclass scaffold NamespacedStateStore in inference
scaffold's state.py promoted the namespace-aware free functions onto
NamespacedStateStore. Update inference accordingly:
- InferenceStateStore now subclasses NamespacedStateStore, gaining
service_state_path / set_current_service / clear_current_service /
current_service_path / resolve_service_name from the parent. It keeps
the inf-specific models_*, lock_model_state, and overrides
recover_pids_from_raw_state to walk both state files.
- ServiceState.save / .remove route through ``store.set_current_service``
/ ``store.clear_current_service`` instead of the deleted free
functions; commands/run.py / register.py / tests resolve log paths
via ``store.logs_dir`` and the service-state path via
``store.service_state_path``.
Behavior unchanged.
* style: sort v2 CLI imports
* refactor(experimental): drop reward verb + RL session flow docs
Address PR areal-project#1434 review:
- Remove ``areal inf reward`` (commands/reward.py + the ``reward_cmd``
wiring in __init__.py) and the corresponding GatewayClient.set_reward
RPC. The reward flow is server-side only for now; the CLI does not
need to wrap it.
- Drop the "RL session flow" and "Setting reward" sections from
cli_guide.md plus the trailing "For plain inference there is no need
to call /rl/start_session" note.
- Drop the ``reward`` entry from the [default].service binding tuple
in config.py and the reward-related cases in the parser test.
---------
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
…les (areal-project#1383) * feat(agent_service): add agent service with OpenClaw and Hermes examples Add an end-to-end agent service stack under areal/v2/agent_service with a gateway bridge, data proxy, and worker, plus two example agents (OpenClaw and Hermes) demonstrating session lifecycle, reward setting, and training integration. Key changes: - Extend gateway bridge and data proxy to drive agent sessions - Add OpenClaw and Hermes example agents under examples/agent_service - Add lifecycle demo, run scripts, and config for the Hermes example - Add integration and per-agent tests under tests/v2/agent_service Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * refactor(examples): drop redundant hermes session self-check scripts The hermes agent service session flow only needs run_agent_service.py: chat sessions are keyed implicitly by the /v1/responses "user" field, and self-evolution mints the per-session sk-sess-* key internally. The start_session.py / demo_lifecycle.py scripts only probed the inference gateway control plane and duplicated examples/openclaw, so remove them. Key changes: - Delete examples/agent_service/hermes/{start_session,demo_lifecycle}.py - Remove the README "Connectivity self-checks" section and Files rows - Repoint set_reward.py docs to run_agent_service.py for the session key * refactor(examples): hoist hermes example to examples/hermes top level Move the Hermes agent-service example from examples/agent_service/hermes to examples/hermes, putting it at the same level as examples/openclaw, and drop the superseded examples/agent_service/openclaw (the standalone examples/openclaw already replaces it). Key changes: - git mv examples/agent_service/hermes -> examples/hermes - Update agent_cls_path and docs to examples.hermes.hermes.HermesAgent - Rewrite README paths from examples/agent_service/hermes to examples/hermes - Remove obsolete examples/agent_service/openclaw Refs: areal-project#1383 Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * refactor(examples): split hermes flow into standalone step scripts Replace the monolithic run_agent_service.py with a 5-step flow matching the openclaw example: train.py, `areal agent run`, start_session.py, hermes_loop.py interaction, and set_reward.py. Update start_session.py for the v2 inference service (HTTP 201 + nested session credentials), fold training defaults into config.yaml, and rewrite the README. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * refactor(examples): inline CLI formatting helpers and trim hermes README Remove the shared examples/hermes/_fmt.py module by inlining the helpers it provided into the two scripts that used it, so each script is self-contained. Also condense the README prose and use a placeholder model path. Key changes: - Inline formatting helpers into start_session.py and set_reward.py - Delete examples/hermes/_fmt.py and its README entry - Simplify quick-start prose; use actor.path=/path/to/your_model - Switch config.yaml actor backend to megatron:d1 Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * docs(examples): clarify session key usage in hermes and openclaw READMEs Annotate which placeholders reuse the sk-sess-* key returned by start_session and distinguish them from the upstream LLM credentials. Trim the openclaw README to focus on the RL training flow. Key changes: - Mark hermes --session-api-key / --api-key as the start_session key - Note HERMES_UPSTREAM_* are your own upstream LLM credentials - Use placeholder values instead of sk-... in env exports - Remove the standalone agent-service section from openclaw README Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * style(test): fix import grouping in agent service tests Remove the stray blank line between the pytest import and the examples.* imports so ruff's isort check passes in CI. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
) Add a reuse_train_logp option to prox_logp_method for decoupled PPO. It reuses the training forward-pass logprobs (detached) as the proximal logp, skipping the extra proximal forward pass and its memory/compute cost. This is only valid with ppo_n_minibatches=1: with a single minibatch the training forward still reflects the policy that generated the rollout, so its logprobs equal the proximal policy. With multiple minibatches the weights change between steps, so PPOActorConfig.__post_init__ rejects that combination. Add tests for the enum/constant wiring, skips_forward_pass, and the ppo_n_minibatches validation.
…roject#1457) reuse_train_logp requires ppo_n_minibatches=1 so the training forward pass still reflects the policy that produced the rollout. Previously an invalid combination raised ValueError; instead warn and force ppo_n_minibatches=1, making clear this reduces the PPO update to a single optimizer step. Remove obsolete tests that asserted behavior no longer relevant to this follow-up.
* chore: add AReaL 2.0 report paper * chore: rename paper file to AReaL2.0_report.pdf
Run a configurable post-exit shell command after local, Ray, and Slurm launchers stop jobs on timeout, interrupt, or failure paths. The hook receives LOG_DIR, has a 600-second timeout, and logs failures without interrupting shutdown or recovery. Ensure the hook still runs if launcher shutdown raises. Add the config field, unit tests, and regenerated CLI docs.
…-project#1460) Add Megatron context-parallel output gathering for forward-only paths and plumb vocabulary logits statistics through PPO/DPO/SFT losses. Support per-key reduce groups in StatsTracker so CP-local loss and vocab statistics can reduce across the appropriate DP/CP group without changing the default reduction group for unrelated stats. Expose Megatron/MoE configuration knobs for router fusion, auxiliary-loss-free balancing, router z-loss, FP32 lm_head output, and fused cross entropy. Update Bailing MoE defaults and regenerate CLI docs.
…ion (areal-project#1454) Use actual trajectory group sizes when applying group-level reward and advantage normalization so failed or filtered rollout samples do not cause fixed-size slices to cross prompt groups. Pass trajectory metadata through batched_call instead of injecting a sentinel batch key, rename the Normalization argument to group_sizes, and zero singleton leave-one-out groups because they have no peer baseline. Add tests for variable-size groups, singleton leave-one-out behavior, validation, 2D advantage normalization, and batched_call metadata forwarding.
…e chat template (areal-project#1463) Some chat templates (e.g. GLM-5.1) iterate over tool_call arguments with .items(), which fails when arguments is a JSON string as per the OpenAI API convention. Normalize tool_call arguments to dicts before applying the chat template in concat_prompt_token_ids_with_parent and the completions/responses tokenizer paths.
Add a SWE-bench RL training workflow under examples/swe: - train_swe_rl.py: the RL training entrypoint. - agent.py: the AReaL-SWEAgent rollout workflow that runs SWE-bench agents through AReaL's OpenAI-compatible proxy during training. - preprocessors.py / prefix_matchers.py: message preprocessing and interaction-cache prefix matching for the agent proxy. - filter_function.py: rollout group accept/reject filter. - qwen3_30b_a3b_grpo.yaml: a runnable Qwen3-Coder GRPO example config. - README.md: setup guide covering the AReaL-SWEAgent checkout, the AEnvironment backend, and Claude Code (cc) agent training. Also add the SWE dataset loader.
* fix: fix safe-to-test CI workflow * fix(tests): remove test_hermes_agent.py Companion cleanup to the openclaw removal in the previous commit. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> * fix: fix wu ci test --------- Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
…upport (areal-project#1458) Extend the OpenAI-compatible proxy rollout server with pluggable message preprocessors and a configurable interaction-cache prefix matcher, add Anthropic-to-OpenAI content handling and tool-call streaming, the Qwen-style tool-call parser, and the accompanying client/proxy/tool-call tests.
…oject#1464) After areal-project#1454, a size-1 leave-one-out group uses the sample itself as the baseline and normalizes to zero instead of passing through the raw advantage. Update test_group_size_edge_cases to assert the new behavior.
swe_sft.py is an SFT-only dataset builder that is not used by the SWE-bench RL training example (examples/swe/train_swe_rl.py loads raw problem instances directly) and is referenced only by its own dataset registration. Remove the module and its registration so the RL example does not ship unused SFT data-processing code.
…er state correctly (areal-project#1468) Since the dist_checkpointing refactor in megatron-core v0.14, flattened_range is legacy, non-reconstructable metadata and is no longer supported as a persistable layout in the checkpoint serialization path; the remaining flattened_range code paths were removed upstream in Megatron-LM PR #2126, and ShardedTensor.validate_metadata_integrity() now rejects any ShardedTensor with flattened_range set. The sharded_state_dict API still defaults to fully_sharded_model_space, which emits flattened_range, so saving a checkpoint with optimizer state fails on the pinned megatron-core 0.17.0. Request dp_reshardable sharding instead, matching the optimizer state format upstream now uses by default. Also pass is_loading=True when building the load-side template. Without it a freshly built optimizer skips megatron-core's state pre-allocation, the template only requests "param", and DCP silently drops exp_avg/exp_avg_sq on resume -- training continues with a reset optimizer state at full learning rate, which we observed to cause gradient-norm spikes and entropy collapse within ~70 steps after recovery in a large-scale RL run.
…ing (areal-project#1470) GroupRMSNorm is a plain nn.Module without its own sharded_state_dict, so Megatron's sharded_state_dict_default treats its weight as replicated. Under TP>1 the saved global shape equals the local shape and only one rank's shard survives a save, corrupting DCP recover and DCP-to-HF conversion (only 1/TP of the gate norm remains). Declare dim 0 as the TP axis via make_sharded_tensors_for_checkpoint so the full num_heads_global * head_dim tensor is stored as proper TP shards.
…ss capacity (areal-project#1471) RolloutController manages staleness globally, but workers also applied their own dp-scaled staleness constraint, dividing capacity by dist.get_world_size(). Force train_data_parallel_size=1 at engine initialization (unless a caller explicitly configures a non-None value) so controller-managed workers do not re-scale capacity; an explicit None must not survive either, or workers fall back to the world-size scaling.
b2a18f0 to
43294f4
Compare
…1502) Point Ascend users to the ascend-v1.0.4 branch and refresh the guide to match its Dockerfiles: CANN 9.0.0 images, pre-built stack table (torch_npu, vLLM-Ascend, MindSpeed, Megatron-Bridge), uv-based source install, and a Geometry3K VLM example replacing the GSM8K quickstart. Sync the Chinese translation with the English version.
* docs(readme): add v2.0.0 release news and update online RL doc link
Add a News entry for the v2.0.0 release highlighting the microservice
architecture (training/inference/agent/weight-update services) and the
two agentic RL training examples (Hermes and SWE agent), with a link to
the technical report on arXiv. Also update the Highlights section's
online RL training link to the new areal-ai.io docs domain.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* docs(readme): refine v2.0 news wording
Drop the "(v2.0.0)" tag from the news date and rephrase the examples
line so each example gets a specific, self-describing name: the Hermes
online RL loop, and end-to-end SWE RL training examples.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* docs(readme): refresh v2.0 highlights and swap logo
- add Hermes and SWE agent examples to Agentic RL table
- collapse [2026/04/23] Scaffoldings news under a <details>
- mark Qwen2.5-VL / Qwen3-VL as Megatron-supported
- link 2026 H2 roadmap issue in Future Roadmap
- add CLI Configurations link under Tutorial
- replace assets/figures/logo.png with v2.0 logo
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* docs: add consolidated CLI guide and tidy README sections
- add docs/{en,zh}/best_practices/cli_guide.md consolidating the
training / inference / agent v2 CLI guides; README Tutorial now
points at the EN version
- Agentic RL table: drop "(v2.0)" labels, rename SWE Agent RL to
Coding Agent RL, note AReaL-SWEAgent/Claude Code Agent support
- fold [2026/04/23], [2026/04/18], [2026/03/02] into the
Previous Releases <details>
- swap Q1 roadmap link for Q2 (areal-project#1302) in Future Roadmap
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* docs(cli): migrate v2 per-service guides into consolidated CLI guide
Content of areal/v2/cli/{training,inference,agent}/cli_guide.md is
now maintained solely in docs/{en,zh}/best_practices/cli_guide.md.
Removed the "canonical source" callback that pointed back at the
per-service files.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* docs(readme): fix Asynchronous RL Guide link path
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* docs(readme): remove gitcgr badge link
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* docs(assets): resize logo to 972x1250
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* style: apply mdformat to README and zh cli_guide
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
---------
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
…ard (areal-project#1529) The padded-vs-packed decision in packed_context_parallel_forward was keyed on whether the individual microbatch carries vision tensors, so an image-free microbatch on a VLM took the wrapper-packed [1, total] path. VLM models cannot consume that layout: their internal packing needs a per-sequence 2D attention mask, so mbridge crashes on the missing mask (attention_mask.sum on None) and megatron-bridge silently computes corrupted positions and packing, producing garbage logprobs for the affected microbatch. Key the decision (and the mask selection) on the model type instead, so text-only microbatches take the same padded branch as image-bearing ones; vision kwargs remain gated on microbatch content. Behavior for non-VLM models and for padded-seq models is unchanged. Validated with Qwen3-VL-2B and Qwen2.5-VL-3B: text-only logprobs match the HuggingFace forward after the fix on both bridges (mean abs diff 1.59 -> 0.03 for the silently corrupted megatron-bridge path).
Allow PPO-family actor losses to weight tokens, sequences, prompt groups, or fixed-length response sums without changing the engine contract. Key changes: - Pair each local policy-gradient mean with its matching engine weight - Preserve explicit prompt-group boundaries across microbatch splitting - Keep token_mean and existing backend behavior unchanged - Document and regression-test padded, packed, ragged, and filtered inputs Refs: areal-project#1423
2496822 to
a8f0d5e
Compare
This fork PR mirrors upstream areal-project#1443 and is targeted at
mainso the alternative implementation can be reviewed without the old areal-project#1417 comparison-base conflict.It keeps the same user-facing actor loss aggregation modes as areal-project#1417, but makes the distributed reduction contract explicit in the training-engine API.
areal-project#1417 uses the existing engine shape: the actor returns a local scalar, a matching weight function returns the local denominator owner count, and each engine reconstructs the global objective as:
This branch represents that pairing as
LossReduction. Each term declares its normalizer and whether the loss function returns a local mean or a local numerator:That interface keeps policy out of the engines. FSDP, Megatron, and Archon all do the same thing: compute local normalizers, all-reduce global normalizers on the correct group, and scale the local loss. The actor owns the policy choice for
actor.loss_aggregation.actor.loss_aggregationtoken_meanseq_meanprompt_meanconstantactive responses * loss_aggregation_divisor)loss_aggregation_divisorBuilt-in PPO/SAPO/CISPO actor loss uses
LossReduction.sum(...)when the policy-gradient loss has pure aggregation semantics, so the engine scales the numerator directly instead of reconstructing it from a local mean. Teacher distillation and M2PO stay onLossReduction.mean(...)because their denominator behavior is local or data-dependent. SFT, DPO, reward-model, and critic training also useLossReduction.mean(...), preserving their previous local-normalized calculation order.Megatron keeps actor loss normalizers on the DP+CP group. With context parallelism, each CP rank sees the full-sequence loss after all-gather, so the DP+CP normalizer cancels the duplicated CP gradient contribution.
Compared with the OSS implementations checked while preparing this change: verl centralizes aggregation around global batch metadata, SkyRL pre-scales before worker loss summation, prime-rl keeps only token mean, and slime needs a CP-aware numerator/denominator reducer. AReaL already has an engine boundary that can carry the distributed normalizer, so this branch makes that boundary first-class instead of adding mode branches inside the engines.
The implementation also excludes denominator-empty responses/groups for
seq_mean,prompt_mean, andconstant. Rejection sampling still usesdenom_maskfor the denominator owner decision, so a pre-rejection active unit can remain in the denominator even if its current numerator is empty.Validation on commit
86b48e5c4:uv run pytest -q tests/test_loss_reduction.py tests/test_prompt_mean_loss.py tests/test_partial_group_norm.py tests/test_grouped_rollout_min_valid.py tests/test_cispo_loss.py tests/test_eval_dispatch.pyuv run ruff check areal/api/engine_api.py areal/engine/core/train_engine.py areal/utils/functional/functional.py areal/trainer/ppo/actor.py tests/test_loss_reduction.py tests/test_prompt_mean_loss.py tests/test_partial_group_norm.py tests/test_grouped_rollout_min_valid.py tests/test_cispo_loss.py tests/test_eval_dispatch.pygit diff --check origin/main...HEADuv run python -m compileall -q areal testssource .venv/bin/activate && pre-commit run --all-filesRuntime demo on commit
86b48e5c4usedexamples/math/boba_grpo.pywith Qwen2.5-0.5B-Instruct, an 8-prompt arithmetic JSONL slice, FSDP actord1, and SGLang rolloutd1. This is a controlled learning-signal demo through the normal PPOTrainer/FSDP/SGLang/weight-update path, not a GSM8K convergence benchmark.The linked W&B runs are sanitized public replays of the captured reward series. They log only the reward metrics and public-safe config fields; code upload, git capture, machine metadata, and system stats were disabled. Static plots are attached from the public fork so the proof does not depend only on the W&B UI remaining available. The exact data slice and command template are in the proof README.
token_meanseq_meanprompt_meanconstant(loss_aggregation_divisor=16)Static reward proof, generated from captured run logs with a trailing 5-step moving average:
Per-mode reward screenshots
Summary by CodeRabbit
New Features
Bug Fixes
Documentation