diff --git a/.agents/skills/review-pr/MIGRATION.md b/.agents/skills/review-pr/MIGRATION.md index 473a24376f..8bcb75c3ff 100644 --- a/.agents/skills/review-pr/MIGRATION.md +++ b/.agents/skills/review-pr/MIGRATION.md @@ -104,8 +104,8 @@ ______________________________________________________________________ 1. `areal/experimental/models/archon/**` + `areal/experimental/engine/archon_engine.py` \+ `areal/experimental/engine/archon_checkpoint.py` -> Distributed Runtime + Model Compute & Attention (CRITICAL/HIGH) -1. `areal/experimental/agent_service/**` -> Service Orchestration (HIGH) -1. `areal/experimental/inference_service/**` -> Service Orchestration (HIGH) +1. `areal/v2/agent_service/**` -> Service Orchestration (HIGH) +1. `areal/v2/inference_service/**` -> Service Orchestration (HIGH) 1. `areal/experimental/engine/archon_weight_sync.py` -> Distributed Runtime (CRITICAL) 1. `areal/infra/rpc/**` -> Launcher & Infrastructure + Distributed Runtime (HIGH) 1. `areal/workflow/**`, `areal/trainer/**` -> Workflow & Trainer Contract (HIGH/MEDIUM) diff --git a/.agents/skills/review-pr/references/review-pr-domains-and-signals.md b/.agents/skills/review-pr/references/review-pr-domains-and-signals.md index a835419ef9..0ce5b3c607 100644 --- a/.agents/skills/review-pr/references/review-pr-domains-and-signals.md +++ b/.agents/skills/review-pr/references/review-pr-domains-and-signals.md @@ -52,19 +52,19 @@ ______________________________________________________________________ ## Domain 4: Service Orchestration (HIGH) -| L2 Signal | File Path Pattern | Code Pattern | -| -------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------- | -| `service_routing_dataflow` | `areal/experimental/agent_service/gateway/`, `areal/experimental/agent_service/router/`, `areal/experimental/inference_service/data_proxy/`, `areal/experimental/inference_service/controller/` | `route`, `gateway`, `router`, `DataProxy`, `controller`, `batch` | -| `session_consistency` | `areal/experimental/agent_service/`, `areal/experimental/inference_service/` | `session`, `affinity`, `history`, `state` | +| L2 Signal | File Path Pattern | Code Pattern | +| -------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------- | +| `service_routing_dataflow` | `areal/v2/agent_service/gateway/`, `areal/v2/agent_service/router/`, `areal/v2/inference_service/data_proxy/`, `areal/v2/inference_service/controller/` | `route`, `gateway`, `router`, `DataProxy`, `controller`, `batch` | +| `session_consistency` | `areal/v2/agent_service/`, `areal/v2/inference_service/` | `session`, `affinity`, `history`, `state` | ## Domain 5: Workflow & Trainer Contract (HIGH/MEDIUM) -| L2 Signal | File Path Pattern | Code Pattern | -| -------------------------- | ----------------------------------------------------------------------------------------------- | --------------------------------------------------- | -| `workflow_engine_boundary` | `areal/workflow/`, `areal/trainer/`, `areal/engine/` | `RolloutWorkflow`, `arun_episode`, `agenerate` | -| `dataset_surface` | `areal/dataset/` | `DataLoader`, `IterableDataset`, `get_*_dataset` | -| `async_contract` | `areal/workflow/`, `areal/experimental/agent_service/`, `areal/experimental/inference_service/` | `async def`, `await`, `aiofiles`, `asyncio` | -| `weight_version_contract` | `areal/api/engine_api.py`, `areal/workflow/`, `areal/trainer/` | `WeightUpdateMeta`, `set_version`, `weight version` | +| L2 Signal | File Path Pattern | Code Pattern | +| -------------------------- | --------------------------------------------------------------------------- | --------------------------------------------------- | +| `workflow_engine_boundary` | `areal/workflow/`, `areal/trainer/`, `areal/engine/` | `RolloutWorkflow`, `arun_episode`, `agenerate` | +| `dataset_surface` | `areal/dataset/` | `DataLoader`, `IterableDataset`, `get_*_dataset` | +| `async_contract` | `areal/workflow/`, `areal/v2/agent_service/`, `areal/v2/inference_service/` | `async def`, `await`, `aiofiles`, `asyncio` | +| `weight_version_contract` | `areal/api/engine_api.py`, `areal/workflow/`, `areal/trainer/` | `WeightUpdateMeta`, `set_version`, `weight version` | ## Domain 6: API & Config Compatibility (MEDIUM) @@ -95,12 +95,12 @@ ______________________________________________________________________ ## Domain 9: Launcher & Infrastructure (HIGH/MEDIUM) -| L2 Signal | File Path Pattern | Code Pattern | -| ------------------------- | ---------------------------------------------------------------------- | --------------------------------------------------------- | -| `launcher_resource_match` | `areal/infra/launcher/` | `LaunchConfig`, `RayLauncher`, `SlurmLauncher` | -| `scheduler_contract` | `areal/infra/scheduler/`, `areal/scheduler/` | `Scheduler`, `placement`, `resource` | -| `rpc_transport` | `areal/infra/rpc/`, `areal/experimental/inference_service/data_proxy/` | `RTensor`, `serialize`, `rpc`, `fetch` | -| `runtime_image_config` | `Dockerfile`, `.dockerignore` | `FROM`, `ARG`, `RUN`, `ENV`, `COPY`, `uv sync`, `VARIANT` | +| L2 Signal | File Path Pattern | Code Pattern | +| ------------------------- | ------------------------------------------------------------ | --------------------------------------------------------- | +| `launcher_resource_match` | `areal/infra/launcher/` | `LaunchConfig`, `RayLauncher`, `SlurmLauncher` | +| `scheduler_contract` | `areal/infra/scheduler/`, `areal/scheduler/` | `Scheduler`, `placement`, `resource` | +| `rpc_transport` | `areal/infra/rpc/`, `areal/v2/inference_service/data_proxy/` | `RTensor`, `serialize`, `rpc`, `fetch` | +| `runtime_image_config` | `Dockerfile`, `.dockerignore` | `FROM`, `ARG`, `RUN`, `ENV`, `COPY`, `uv sync`, `VARIANT` | ## Domain 10: Low-Risk Hygiene (LOW) diff --git a/.agents/skills/upgrade-deps/checklists/sglang.md b/.agents/skills/upgrade-deps/checklists/sglang.md index 11c05eb569..f9d34e46d3 100644 --- a/.agents/skills/upgrade-deps/checklists/sglang.md +++ b/.agents/skills/upgrade-deps/checklists/sglang.md @@ -23,27 +23,27 @@ upstream_paths: ### Primary (engine layer — most likely to break) -| File | Imports / Usage | -| -------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `areal/engine/sglang_remote.py` | HTTP-only (no Python imports). Sends requests to `/generate`, `/load_lora_adapter`, `/update_weights_from_disk`, `/update_weights_from_distributed`, `/init_weights_update_group`, `/pause_generation`, `/continue_generation`, `/health`, `/release_memory_occupation`, `/resume_memory_occupation`. Parses `meta_info`, `finish_reason`, `output_token_logprobs`, `routed_experts` from JSON responses. | -| `areal/experimental/openai/tool_call_parser.py` | `sglang.srt.entrypoints.openai.protocol.Function`, `sglang.srt.entrypoints.openai.protocol.Tool`, `sglang.srt.function_call.function_call_parser.FunctionCallParser`, `sglang.srt.parser.reasoning_parser.ReasoningParser` | -| `areal/experimental/inference_service/sglang/launch_server.py` | `sglang.srt.entrypoints.engine.Engine`, `sglang.srt.entrypoints.engine.init_tokenizer_manager`, `sglang.srt.entrypoints.http_server._execute_server_warmup`, `sglang.srt.entrypoints.http_server._setup_and_run_http_server`, `sglang.srt.entrypoints.http_server.app`, `sglang.srt.managers.detokenizer_manager.run_detokenizer_process` | -| `areal/experimental/inference_service/sglang/scheduler.py` | `sglang.srt.managers.scheduler.Scheduler`, `sglang.srt.managers.scheduler.configure_scheduler`, `sglang.srt.observability.trace.process_tracing_init`, `sglang.srt.observability.trace.trace_set_thread_info`, `sglang.srt.utils.get_bool_env_var`, `sglang.srt.utils.kill_itself_when_parent_died`, `sglang.srt.utils.set_gpu_proc_affinity`, `sglang.srt.utils.numa_utils.get_numa_node_if_available`, `sglang.srt.utils.numa_utils.numa_bind_to_node`, `sglang.srt.environ.envs`, `sglang.utils.get_exception_traceback` | -| `areal/experimental/inference_service/sglang/rpc_proxy.py` | `sglang.srt.managers.io_struct.RpcReqInput`, `sglang.srt.managers.io_struct.RpcReqOutput`, `sglang.srt.server_args.PortArgs`, `sglang.srt.utils.network.get_zmq_socket` | -| `areal/experimental/inference_service/sglang/bridge.py` | HTTP-only (same protocol as `sglang_remote.py`). `/generate` request/response, `/pause_generation`, `/continue_generation`, `/release_memory_occupation`, `/resume_memory_occupation`. | +| File | Imports / Usage | +| ---------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `areal/engine/sglang_remote.py` | HTTP-only (no Python imports). Sends requests to `/generate`, `/load_lora_adapter`, `/update_weights_from_disk`, `/update_weights_from_distributed`, `/init_weights_update_group`, `/pause_generation`, `/continue_generation`, `/health`, `/release_memory_occupation`, `/resume_memory_occupation`. Parses `meta_info`, `finish_reason`, `output_token_logprobs`, `routed_experts` from JSON responses. | +| `areal/experimental/openai/tool_call_parser.py` | `sglang.srt.entrypoints.openai.protocol.Function`, `sglang.srt.entrypoints.openai.protocol.Tool`, `sglang.srt.function_call.function_call_parser.FunctionCallParser`, `sglang.srt.parser.reasoning_parser.ReasoningParser` | +| `areal/v2/inference_service/sglang/launch_server.py` | `sglang.srt.entrypoints.engine.Engine`, `sglang.srt.entrypoints.engine.init_tokenizer_manager`, `sglang.srt.entrypoints.http_server._execute_server_warmup`, `sglang.srt.entrypoints.http_server._setup_and_run_http_server`, `sglang.srt.entrypoints.http_server.app`, `sglang.srt.managers.detokenizer_manager.run_detokenizer_process` | +| `areal/v2/inference_service/sglang/scheduler.py` | `sglang.srt.managers.scheduler.Scheduler`, `sglang.srt.managers.scheduler.configure_scheduler`, `sglang.srt.observability.trace.process_tracing_init`, `sglang.srt.observability.trace.trace_set_thread_info`, `sglang.srt.utils.get_bool_env_var`, `sglang.srt.utils.kill_itself_when_parent_died`, `sglang.srt.utils.set_gpu_proc_affinity`, `sglang.srt.utils.numa_utils.get_numa_node_if_available`, `sglang.srt.utils.numa_utils.numa_bind_to_node`, `sglang.srt.environ.envs`, `sglang.utils.get_exception_traceback` | +| `areal/v2/inference_service/sglang/rpc_proxy.py` | `sglang.srt.managers.io_struct.RpcReqInput`, `sglang.srt.managers.io_struct.RpcReqOutput`, `sglang.srt.server_args.PortArgs`, `sglang.srt.utils.network.get_zmq_socket` | +| `areal/v2/inference_service/sglang/bridge.py` | HTTP-only (same protocol as `sglang_remote.py`). `/generate` request/response, `/pause_generation`, `/continue_generation`, `/release_memory_occupation`, `/resume_memory_occupation`. | ### Secondary (model / infra layer) -| File | Imports / Usage | -| --------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------ | -| `areal/infra/launcher/sglang_server.py` | HTTP-only. Polls `/v1/models` to confirm server readiness. Calls `SGLangConfig.build_cmd()` to construct the server launch command. | -| `areal/api/cli_args.py` | `SGLangConfig.build_cmd()`, `SGLangConfig.build_cmd_from_args()`, `SGLangConfig.build_args()`. Version check against `0.5.10.post1`. | -| `areal/trainer/rl_trainer.py` | `SGLangConfig.build_args()` to construct `RemoteSGLangEngine`. | -| `areal/infra/launcher/ray.py` | `SGLangConfig` import + `to_structured_cfg`. | -| `areal/infra/launcher/local.py` | `SGLangConfig` import + `to_structured_cfg`. | -| `areal/infra/launcher/slurm.py` | `SGLangConfig` import + `to_structured_cfg`. | -| `areal/experimental/inference_service/controller/controller.py` | `SGLangConfig` in gateway controller. | -| `areal/experimental/inference_service/data_proxy/backend.py` | HTTP protocol implementation for inference service proxy (SGLangBridgeBackend). | +| File | Imports / Usage | +| ----------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------ | +| `areal/infra/launcher/sglang_server.py` | HTTP-only. Polls `/v1/models` to confirm server readiness. Calls `SGLangConfig.build_cmd()` to construct the server launch command. | +| `areal/api/cli_args.py` | `SGLangConfig.build_cmd()`, `SGLangConfig.build_cmd_from_args()`, `SGLangConfig.build_args()`. Version check against `0.5.10.post1`. | +| `areal/trainer/rl_trainer.py` | `SGLangConfig.build_args()` to construct `RemoteSGLangEngine`. | +| `areal/infra/launcher/ray.py` | `SGLangConfig` import + `to_structured_cfg`. | +| `areal/infra/launcher/local.py` | `SGLangConfig` import + `to_structured_cfg`. | +| `areal/infra/launcher/slurm.py` | `SGLangConfig` import + `to_structured_cfg`. | +| `areal/v2/inference_service/controller/controller.py` | `SGLangConfig` in gateway controller. | +| `areal/v2/inference_service/data_proxy/backend.py` | HTTP protocol implementation for inference service proxy (SGLangBridgeBackend). | ### Tertiary (tests, config) @@ -415,8 +415,7 @@ ______________________________________________________________________ **Source:** `python/sglang/srt/entrypoints/engine.py` -Called in `areal/experimental/inference_service/sglang/launch_server.py` -(`areal_launch_server`): +Called in `areal/v2/inference_service/sglang/launch_server.py` (`areal_launch_server`): ```python from sglang.srt.entrypoints.engine import Engine, init_tokenizer_manager @@ -447,8 +446,7 @@ ______________________________________________________________________ **Source:** `python/sglang/srt/entrypoints/http_server.py` -Called in `areal/experimental/inference_service/sglang/launch_server.py` -(`areal_launch_server`): +Called in `areal/v2/inference_service/sglang/launch_server.py` (`areal_launch_server`): ```python from sglang.srt.entrypoints.http_server import ( @@ -478,7 +476,7 @@ ______________________________________________________________________ **Source:** `python/sglang/srt/managers/scheduler.py` -Called in `areal/experimental/inference_service/sglang/scheduler.py` +Called in `areal/v2/inference_service/sglang/scheduler.py` (`areal_run_scheduler_process`): ```python @@ -510,7 +508,7 @@ ______________________________________________________________________ **Source:** `python/sglang/srt/managers/io_struct.py` -Used in `areal/experimental/inference_service/sglang/rpc_proxy.py`: +Used in `areal/v2/inference_service/sglang/rpc_proxy.py`: ```python from sglang.srt.managers.io_struct import RpcReqInput, RpcReqOutput diff --git a/.claude/data/review-pr-domains-and-signals.md b/.claude/data/review-pr-domains-and-signals.md index 7587d6dda9..16d71ff7ac 100644 --- a/.claude/data/review-pr-domains-and-signals.md +++ b/.claude/data/review-pr-domains-and-signals.md @@ -52,19 +52,19 @@ ______________________________________________________________________ ## Domain 4: Service Orchestration (HIGH) -| L2 Signal | File Path Pattern | Code Pattern | -| -------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------- | -| `service_routing_dataflow` | `areal/experimental/agent_service/gateway/`, `areal/experimental/agent_service/router/`, `areal/experimental/inference_service/data_proxy/`, `areal/experimental/inference_service/controller/` | `route`, `gateway`, `router`, `DataProxy`, `controller`, `batch` | -| `session_consistency` | `areal/experimental/agent_service/`, `areal/experimental/inference_service/` | `session`, `affinity`, `history`, `state` | +| L2 Signal | File Path Pattern | Code Pattern | +| -------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------- | +| `service_routing_dataflow` | `areal/v2/agent_service/gateway/`, `areal/v2/agent_service/router/`, `areal/v2/inference_service/data_proxy/`, `areal/v2/inference_service/controller/` | `route`, `gateway`, `router`, `DataProxy`, `controller`, `batch` | +| `session_consistency` | `areal/v2/agent_service/`, `areal/v2/inference_service/` | `session`, `affinity`, `history`, `state` | ## Domain 5: Workflow & Trainer Contract (HIGH/MEDIUM) -| L2 Signal | File Path Pattern | Code Pattern | -| -------------------------- | ----------------------------------------------------------------------------------------------- | --------------------------------------------------- | -| `workflow_engine_boundary` | `areal/workflow/`, `areal/trainer/`, `areal/engine/` | `RolloutWorkflow`, `arun_episode`, `agenerate` | -| `dataset_surface` | `areal/dataset/` | `DataLoader`, `IterableDataset`, `get_*_dataset` | -| `async_contract` | `areal/workflow/`, `areal/experimental/agent_service/`, `areal/experimental/inference_service/` | `async def`, `await`, `aiofiles`, `asyncio` | -| `weight_version_contract` | `areal/api/engine_api.py`, `areal/workflow/`, `areal/trainer/` | `WeightUpdateMeta`, `set_version`, `weight version` | +| L2 Signal | File Path Pattern | Code Pattern | +| -------------------------- | --------------------------------------------------------------------------- | --------------------------------------------------- | +| `workflow_engine_boundary` | `areal/workflow/`, `areal/trainer/`, `areal/engine/` | `RolloutWorkflow`, `arun_episode`, `agenerate` | +| `dataset_surface` | `areal/dataset/` | `DataLoader`, `IterableDataset`, `get_*_dataset` | +| `async_contract` | `areal/workflow/`, `areal/v2/agent_service/`, `areal/v2/inference_service/` | `async def`, `await`, `aiofiles`, `asyncio` | +| `weight_version_contract` | `areal/api/engine_api.py`, `areal/workflow/`, `areal/trainer/` | `WeightUpdateMeta`, `set_version`, `weight version` | ## Domain 6: API & Config Compatibility (MEDIUM) @@ -95,12 +95,12 @@ ______________________________________________________________________ ## Domain 9: Launcher & Infrastructure (HIGH/MEDIUM) -| L2 Signal | File Path Pattern | Code Pattern | -| ------------------------- | ---------------------------------------------------------------------- | --------------------------------------------------------- | -| `launcher_resource_match` | `areal/infra/launcher/` | `LaunchConfig`, `RayLauncher`, `SlurmLauncher` | -| `scheduler_contract` | `areal/infra/scheduler/`, `areal/scheduler/` | `Scheduler`, `placement`, `resource` | -| `rpc_transport` | `areal/infra/rpc/`, `areal/experimental/inference_service/data_proxy/` | `RTensor`, `serialize`, `rpc`, `fetch` | -| `runtime_image_config` | `Dockerfile`, `.dockerignore` | `FROM`, `ARG`, `RUN`, `ENV`, `COPY`, `uv sync`, `VARIANT` | +| L2 Signal | File Path Pattern | Code Pattern | +| ------------------------- | ------------------------------------------------------------ | --------------------------------------------------------- | +| `launcher_resource_match` | `areal/infra/launcher/` | `LaunchConfig`, `RayLauncher`, `SlurmLauncher` | +| `scheduler_contract` | `areal/infra/scheduler/`, `areal/scheduler/` | `Scheduler`, `placement`, `resource` | +| `rpc_transport` | `areal/infra/rpc/`, `areal/v2/inference_service/data_proxy/` | `RTensor`, `serialize`, `rpc`, `fetch` | +| `runtime_image_config` | `Dockerfile`, `.dockerignore` | `FROM`, `ARG`, `RUN`, `ENV`, `COPY`, `uv sync`, `VARIANT` | ## Domain 10: Low-Risk Hygiene (LOW) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index cd1300bc62..7b5cd04f58 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -1,34 +1,34 @@ # Default: catch-all (lead + two active maintainers as backups). # Last matching rule wins, so this is the fallback for any path # not covered by a more specific rule below. -* @garrett4wade @sitabulaixizawaluduo @fishcrap +* @garrett4wade @sitabulaixizawaluduo @fishcrap # Core package -/areal/api/ @garrett4wade @rchardx @sitabulaixizawaluduo -/areal/engine/ @rchardx @nuzant @garrett4wade @geshi001 -/areal/experimental/inference_service/ @nuzant @guozhihao-224 @TaoZex -/areal/experimental/agent_service/ @CormickKneey @nuzant -/areal/experimental/training_service/ @sitabulaixizawaluduo @garrett4wade -/areal/experimental/weight_update/ @TaoZex @garrett4wade @sitabulaixizawaluduo -/areal/infra/ @HwVanICI @garrett4wade @guozhihao-224 -/areal/models/ @rchardx @nuzant @geshi001 @PrometheusComing -/areal/trainer/ @garrett4wade @rchardx @fishcrap +/areal/api/ @garrett4wade @rchardx @sitabulaixizawaluduo +/areal/engine/ @rchardx @nuzant @garrett4wade @geshi001 +/areal/v2/inference_service/ @nuzant @guozhihao-224 @TaoZex +/areal/v2/agent_service/ @CormickKneey @nuzant +/areal/v2/training_service/ @sitabulaixizawaluduo @garrett4wade +/areal/v2/weight_update/ @TaoZex @garrett4wade @sitabulaixizawaluduo +/areal/infra/ @HwVanICI @garrett4wade @guozhihao-224 +/areal/models/ @rchardx @nuzant @geshi001 @PrometheusComing +/areal/trainer/ @garrett4wade @rchardx @fishcrap # Tests & Examples -/tests/ @garrett4wade @sitabulaixizawaluduo -/examples/ @zhenanf @HwVanICI @CormickKneey @PrometheusComing +/tests/ @garrett4wade @sitabulaixizawaluduo +/examples/ @zhenanf @HwVanICI @CormickKneey @PrometheusComing # Documentation -/docs/ @garrett4wade @nuzant +/docs/ @garrett4wade @nuzant # CI/CD & infrastructure -/.github/ @garrett4wade @nuzant @sitabulaixizawaluduo -/Dockerfile @garrett4wade @fishcrap @sitabulaixizawaluduo -pyproject.toml @garrett4wade @fishcrap @sitabulaixizawaluduo -pyproject.vllm.toml @garrett4wade @fishcrap @sitabulaixizawaluduo -uv.lock @garrett4wade @fishcrap @sitabulaixizawaluduo -uv.vllm.lock @garrett4wade @fishcrap @sitabulaixizawaluduo +/.github/ @garrett4wade @nuzant @sitabulaixizawaluduo +/Dockerfile @garrett4wade @fishcrap @sitabulaixizawaluduo +pyproject.toml @garrett4wade @fishcrap @sitabulaixizawaluduo +pyproject.vllm.toml @garrett4wade @fishcrap @sitabulaixizawaluduo +uv.lock @garrett4wade @fishcrap @sitabulaixizawaluduo +uv.vllm.lock @garrett4wade @fishcrap @sitabulaixizawaluduo # Governance & community -GOVERNANCE.md @garrett4wade @sitabulaixizawaluduo -CONTRIBUTING.md @garrett4wade @sitabulaixizawaluduo +GOVERNANCE.md @garrett4wade @sitabulaixizawaluduo +CONTRIBUTING.md @garrett4wade @sitabulaixizawaluduo diff --git a/.github/workflows/tag-release-image.yml b/.github/workflows/tag-release-image.yml index 0c3589d8b7..18651a4bc2 100644 --- a/.github/workflows/tag-release-image.yml +++ b/.github/workflows/tag-release-image.yml @@ -6,7 +6,7 @@ on: workflow_dispatch: inputs: tag: - description: 'Release tag to build (e.g., v1.0.4-sglang)' + description: 'Release tag to build (e.g., v2.0.0-sglang)' required: true type: string diff --git a/.github/workflows/test-areal.yml b/.github/workflows/test-areal.yml index f38d1a01fb..01357a8deb 100644 --- a/.github/workflows/test-areal.yml +++ b/.github/workflows/test-areal.yml @@ -322,7 +322,7 @@ jobs: VIRTUAL_ENV: /opt/.venv run: | export PATH="/opt/.venv/bin:$PATH" - pytest -m "(not slow or ci) and not ${EXCLUDE_BACKEND}" --durations=20 -s -vv tests/test_*.py tests/experimental/ tests/infra/ + pytest -m "(not slow or ci) and not ${EXCLUDE_BACKEND}" --durations=20 -s -vv tests/test_*.py tests/experimental/ tests/infra/ tests/v2/ cleanup: name: Tear down GCP runner (${{ matrix.variant }}) diff --git a/.opencode/data/review-pr-domains-and-signals.md b/.opencode/data/review-pr-domains-and-signals.md index 21c65a15af..fe9d9aef17 100644 --- a/.opencode/data/review-pr-domains-and-signals.md +++ b/.opencode/data/review-pr-domains-and-signals.md @@ -52,19 +52,19 @@ ______________________________________________________________________ ## Domain 4: Service Orchestration (HIGH) -| L2 Signal | File Path Pattern | Code Pattern | -| -------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------- | -| `service_routing_dataflow` | `areal/experimental/agent_service/gateway/`, `areal/experimental/agent_service/router/`, `areal/experimental/inference_service/data_proxy/`, `areal/experimental/inference_service/controller/` | `route`, `gateway`, `router`, `DataProxy`, `controller`, `batch` | -| `session_consistency` | `areal/experimental/agent_service/`, `areal/experimental/inference_service/` | `session`, `affinity`, `history`, `state` | +| L2 Signal | File Path Pattern | Code Pattern | +| -------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------- | +| `service_routing_dataflow` | `areal/v2/agent_service/gateway/`, `areal/v2/agent_service/router/`, `areal/v2/inference_service/data_proxy/`, `areal/v2/inference_service/controller/` | `route`, `gateway`, `router`, `DataProxy`, `controller`, `batch` | +| `session_consistency` | `areal/v2/agent_service/`, `areal/v2/inference_service/` | `session`, `affinity`, `history`, `state` | ## Domain 5: Workflow & Trainer Contract (HIGH/MEDIUM) -| L2 Signal | File Path Pattern | Code Pattern | -| -------------------------- | ----------------------------------------------------------------------------------------------- | --------------------------------------------------- | -| `workflow_engine_boundary` | `areal/workflow/`, `areal/trainer/`, `areal/engine/` | `RolloutWorkflow`, `arun_episode`, `agenerate` | -| `dataset_surface` | `areal/dataset/` | `DataLoader`, `IterableDataset`, `get_*_dataset` | -| `async_contract` | `areal/workflow/`, `areal/experimental/agent_service/`, `areal/experimental/inference_service/` | `async def`, `await`, `aiofiles`, `asyncio` | -| `weight_version_contract` | `areal/api/engine_api.py`, `areal/workflow/`, `areal/trainer/` | `WeightUpdateMeta`, `set_version`, `weight version` | +| L2 Signal | File Path Pattern | Code Pattern | +| -------------------------- | --------------------------------------------------------------------------- | --------------------------------------------------- | +| `workflow_engine_boundary` | `areal/workflow/`, `areal/trainer/`, `areal/engine/` | `RolloutWorkflow`, `arun_episode`, `agenerate` | +| `dataset_surface` | `areal/dataset/` | `DataLoader`, `IterableDataset`, `get_*_dataset` | +| `async_contract` | `areal/workflow/`, `areal/v2/agent_service/`, `areal/v2/inference_service/` | `async def`, `await`, `aiofiles`, `asyncio` | +| `weight_version_contract` | `areal/api/engine_api.py`, `areal/workflow/`, `areal/trainer/` | `WeightUpdateMeta`, `set_version`, `weight version` | ## Domain 6: API & Config Compatibility (MEDIUM) @@ -95,12 +95,12 @@ ______________________________________________________________________ ## Domain 9: Launcher & Infrastructure (HIGH/MEDIUM) -| L2 Signal | File Path Pattern | Code Pattern | -| ------------------------- | ---------------------------------------------------------------------- | --------------------------------------------------------- | -| `launcher_resource_match` | `areal/infra/launcher/` | `LaunchConfig`, `RayLauncher`, `SlurmLauncher` | -| `scheduler_contract` | `areal/infra/scheduler/`, `areal/scheduler/` | `Scheduler`, `placement`, `resource` | -| `rpc_transport` | `areal/infra/rpc/`, `areal/experimental/inference_service/data_proxy/` | `RTensor`, `serialize`, `rpc`, `fetch` | -| `runtime_image_config` | `Dockerfile`, `.dockerignore` | `FROM`, `ARG`, `RUN`, `ENV`, `COPY`, `uv sync`, `VARIANT` | +| L2 Signal | File Path Pattern | Code Pattern | +| ------------------------- | ------------------------------------------------------------ | --------------------------------------------------------- | +| `launcher_resource_match` | `areal/infra/launcher/` | `LaunchConfig`, `RayLauncher`, `SlurmLauncher` | +| `scheduler_contract` | `areal/infra/scheduler/`, `areal/scheduler/` | `Scheduler`, `placement`, `resource` | +| `rpc_transport` | `areal/infra/rpc/`, `areal/v2/inference_service/data_proxy/` | `RTensor`, `serialize`, `rpc`, `fetch` | +| `runtime_image_config` | `Dockerfile`, `.dockerignore` | `FROM`, `ARG`, `RUN`, `ENV`, `COPY`, `uv sync`, `VARIANT` | ## Domain 10: Low-Risk Hygiene (LOW) diff --git a/Dockerfile b/Dockerfile index 7c3e1f8239..5d2f6f5fe6 100644 --- a/Dockerfile +++ b/Dockerfile @@ -264,7 +264,7 @@ RUN set -ex \ ENV PATH="/root/.cargo/bin:$PATH" RUN curl -fsSL \ https://github.com/zeroclaw-labs/zeroclaw/releases/latest/download/install.sh \ - | bash -s -- --skip-onboard \ + | bash -s -- --skip-quickstart \ && rm -rf /tmp/* ############################################################## diff --git a/README.md b/README.md index 453d2a8391..9b20e92d01 100644 --- a/README.md +++ b/README.md @@ -41,6 +41,12 @@ cost-effective** for a broad community of developers and researchers. ## 📰 News +**\[2026/06/17\]** 🔬 Introducing **KPop** — bidirectional binary KL divergence token +masking. Configured via `rejection_sampling.metric=binary_kl`. Also adding an +[IcePop config](examples/math/gsm8k_icepop.yaml) (importance-ratio-based token masking). +Check out [gsm8k_kpop.yaml](examples/math/gsm8k_kpop.yaml) and +[gsm8k_icepop.yaml](examples/math/gsm8k_icepop.yaml) to get started! + **\[2026/04/23\]** 🚀 We’re excited to release our integration with [Scaffoldings](https://github.com/NVIDIA/TensorRT-LLM/tree/main/tensorrt_llm/scaffolding) for agentic RL training - now live in our @@ -207,6 +213,8 @@ All RL algorithms support both asynchronous and synchronous versions by setting | **REINFORCE++** | - | [📄 Paper](https://arxiv.org/pdf/2501.03262) | [🔗 GSM8K Example](examples/math/gsm8k_reinforce.yaml) | | **RLOO** | [📖 Docs](docs/en/algorithms/grpo_series.md) | [📄 Paper](https://arxiv.org/pdf/2402.14740v1) | [🔗 GSM8K Example](examples/math/gsm8k_rloo.yaml) | | **SAPO** | [📖 Docs](docs/en/algorithms/grpo_series.md) | [📄 Paper](https://arxiv.org/abs/2511.20347) | [🔗 GSM8K Example](examples/math/gsm8k_sapo.yaml) | +| **IcePop** | [📖 Docs](docs/en/algorithms/grpo_series.md) | [📄 Blog](https://ringtech.notion.site/icepop) | [🔗 GSM8K Example](examples/math/gsm8k_icepop.yaml) | +| **KPop** | [📖 Docs](docs/en/algorithms/grpo_series.md) | [📄 Blog](https://ringtech.notion.site/kpop) | [🔗 GSM8K Example](examples/math/gsm8k_kpop.yaml) | | **M2PO** | [📖 Docs](docs/algorithms/m2po.md) | [📄 Paper](https://arxiv.org/abs/2510.01161) | [🔗 GSM8K Example](examples/math/gsm8k_m2po.yaml) | | **DPO** | [📖 Docs](docs/en/algorithms/dpo.md) | [📄 Paper](https://arxiv.org/abs/2305.18290) | [🔗 HH-RLHF Example](examples/alignment/hhrlhf_dpo.yaml) | | **RLHF Reward Modeling** | - | - | [🔗 RLHF Example](examples/alignment/hhrlhf_rw.yaml) | @@ -332,6 +340,18 @@ monthly. We warmly welcome community engagement and contributions. We are also **actively hiring interns and full-time employees** with open positions in both the US and China. +## 📢 Community + +WeChat QR Code + +We warmly welcome everyone to join our community! We have a _**WeChat group**_ for +real-time discussions and a +[GitHub discussion board](https://github.com/areal-project/AReaL/discussions) for +asynchronous conversations. + +Check out our [Community Repository](https://github.com/areal-project/community) for +more details and meeting materials. + ## 🙏 Acknowledgments We gratefully acknowledge that major contributors are from the AReaL Team at the diff --git a/ROADMAP.md b/ROADMAP.md index 09c64d94af..68c525aded 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -91,9 +91,9 @@ in contributing, please reach out to discuss implementation details. - [ ] Migrate primary Megatron integration lib from `mbridge` to `megatron-bridge` ([#1260](https://github.com/areal-project/AReaL/issues/1260)) `[CC]` - [ ] Full support for colocation/separation weight transfer with `awex` as the backend - (`areal/experimental/weight_update/`) `[CC]` + (`areal/v2/weight_update/`) `[CC]` - [ ] Migrate legacy NCCL broadcast weight transfer approach from `areal/engine` into - `areal/experimental/weight_update/` `[CC]` + `areal/v2/weight_update/` `[CC]` - [ ] Omini model RL support with FSDP backend ([#879](https://github.com/areal-project/AReaL/issues/879)) `[CC]` - [ ] Support training the latest large MoE models with the Megatron backend, including diff --git a/areal/api/__init__.py b/areal/api/__init__.py index 3401bf8b89..b7e7316776 100644 --- a/areal/api/__init__.py +++ b/areal/api/__init__.py @@ -3,6 +3,8 @@ __all__ = [ "RolloutWorkflow", "AsyncRewardWrapper", + "RewardResult", + "normalize_reward_result", "TrainEngine", "InferenceEngine", "Scheduler", @@ -49,6 +51,8 @@ "WorkflowLike": "areal.api.workflow_api", "AgentWorkflow": "areal.api.workflow_api", "AsyncRewardWrapper": "areal.api.reward_api", + "RewardResult": "areal.api.reward_api", + "normalize_reward_result": "areal.api.reward_api", "RolloutWorkflow": "areal.api.workflow_api", } diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py index d5ae9c2cad..4ddcf613ac 100644 --- a/areal/api/cli_args.py +++ b/areal/api/cli_args.py @@ -921,6 +921,45 @@ class MegatronEngineConfig: default=False, metadata={"help": "Fuse token rearrangement ops during token dispatching."}, ) + moe_router_fusion: bool = field( + default=False, + metadata={ + "help": "Enable fusion for MoE TopK routing and aux-loss computation. " + "Requires TransformerEngine >= 2.7.0.", + }, + ) + moe_router_bias_update_rate: float = field( + default=0.0, + metadata={ + "help": "Update rate for auxiliary-loss-free MoE load balancing " + "(DeepSeek V3 style). Controls how fast expert_bias adjusts. " + "Default 0.0 disables bias updates; set a positive value such as " + "1e-3 to enable.", + }, + ) + moe_z_loss_coeff: float | None = field( + default=None, + metadata={ + "help": "Scaling coefficient for router z-loss. Complements " + "auxiliary-loss-free load balancing for router stability. A starting " + "value of 1e-3 is recommended. None disables z-loss.", + }, + ) + + # Precision & Loss + enable_fp32_lm_head: bool = field( + default=False, + metadata={ + "help": "Cast lm_head output to FP32 before loss computation for " + "numerical stability." + }, + ) + cross_entropy_loss_fusion: bool = field( + default=False, + metadata={ + "help": "Enable fused cross-entropy loss kernel for better performance." + }, + ) # FP8 Training Configuration fp8_config: FP8EngineConfig | None = None @@ -951,6 +990,24 @@ class MegatronEngineConfig: }, ) + disable_grad_buffers_cpu_backup: bool = field( + default=False, + metadata={ + "help": ( + "When offloading with torch_memory_saver, skip CPU backup for " + "Megatron gradient buffers (they are recomputed each step). " + ) + }, + ) + + enable_mtp: bool = field( + default=False, + metadata={ + "help": "Keep the model's Multi-Token-Prediction (MTP) head " + "(bridge_type=megatron-bridge only). Default False drops it.", + }, + ) + class SchedulingStrategyType(str, Enum): separation = "separation" @@ -1348,8 +1405,10 @@ class RejectionSamplingConfig: "'ratio': direct importance ratio π_proximal/π_behave. " "'kl_k1': KL estimator k1 = log(r), forward KL unbiased estimator (can be negative). " "'kl_k2': KL estimator k2 = 0.5 * (log r)^2, non-negative quadratic approximation. " - "'kl_k3': KL estimator k3 = r - log(r) - 1, non-negative exact forward KL estimator.", - "choices": ["ratio", "kl_k1", "kl_k2", "kl_k3"], + "'kl_k3': KL estimator k3 = r - log(r) - 1, non-negative exact forward KL estimator. " + "'binary_kl': KPop (symmetric binary KL divergence) — masks tokens where either " + "KL(proximal||behave) or KL(behave||proximal) exceeds the upper bound.", + "choices": ["ratio", "kl_k1", "kl_k2", "kl_k3", "binary_kl"], }, ) agg: str = field( @@ -1545,6 +1604,18 @@ class PPOActorConfig(TrainEngineConfig): metadata={"help": "SAPO temperature for negative advantages"}, ) + # CISPO (Clipped IS-weight Policy Optimization) - MiniMax-M1 https://arxiv.org/abs/2506.13585 + use_cispo_loss: bool = field( + default=False, + metadata={ + "help": "Use CISPO loss: clip the importance-sampling weight under " + "stop-gradient and keep gradient on every token's log pi (MiniMax-M1 " + "Eq. 4-5). Mutually exclusive with SAPO. Token-level only. Requires " + "eps_clip_higher > 0; recommended eps_clip=1.0 (single-sided, lower " + "bound 0) with eps_clip_higher=4.0." + }, + ) + # Asynchronous RL recompute_logprob: bool = field( default=False, @@ -1581,7 +1652,9 @@ class PPOActorConfig(TrainEngineConfig): "Only effective when use_decoupled_loss=True. Options: " "'recompute' (default): Standard decoupled PPO, recompute proximal policy via forward pass. " "'loglinear': Use log-linear interpolation to approximate proximal policy (skip forward pass). " - "'metrics': Like 'recompute', but also compute approximation metrics for evaluation.", + "'metrics': Like 'recompute', but also compute approximation metrics for evaluation. " + "'reuse_train_logp': Reuse training forward-pass logprobs as the proximal " + "logp (skip the extra forward; requires ppo_n_minibatches=1).", "choices": PROX_LOGP_METHODS_ALL, }, ) @@ -1616,6 +1689,21 @@ def should_compute_prox_logp(self) -> bool: def __post_init__(self): """Validate PPO actor configuration.""" + from areal.utils.constants import ProxLogpMethod + + if ( + ProxLogpMethod(self.prox_logp_method) == ProxLogpMethod.REUSE_TRAIN_LOGP + and self.ppo_n_minibatches > 1 + ): + logger.warning( + "prox_logp_method='reuse_train_logp' requires ppo_n_minibatches=1, " + f"but got ppo_n_minibatches={self.ppo_n_minibatches}. " + "With multiple minibatches, weights change between steps, so the " + "training forward logprobs would differ from the original policy. " + "Forcing ppo_n_minibatches=1; note this changes training dynamics " + "to a single optimizer step per PPO update." + ) + self.ppo_n_minibatches = 1 # Warn if rejection_sampling is configured but use_decoupled_loss is False if not self.use_decoupled_loss and self.rejection_sampling is not None: logger.warning( @@ -1652,6 +1740,25 @@ def __post_init__(self): "Please set `actor.use_decoupled_loss=false` in your configuration." ) + # Validate CISPO configuration + if self.use_cispo_loss: + if self.use_sapo_loss: + raise ValueError( + "CISPO and SAPO are mutually exclusive surrogates. " + "Set at most one of use_cispo_loss / use_sapo_loss." + ) + if self.eps_clip_higher is None or self.eps_clip_higher <= 0: + raise ValueError( + "CISPO requires a positive eps_clip_higher (the asymmetric " + "upper clip is its defining knob, MiniMax-M1 Eq. 4-5). Got " + f"eps_clip_higher={self.eps_clip_higher}." + ) + if self.importance_sampling_level != "token": + raise ValueError( + "CISPO only supports importance_sampling_level='token'. " + "Sequence-level (GSPO-style) CISPO has no published surrogate." + ) + super().__post_init__() @@ -1914,9 +2021,7 @@ def build_cmd( @staticmethod def build_cmd_from_args(args: dict[str, Any]): - return get_py_cmd( - "areal.experimental.inference_service.sglang.launch_server", args - ) + return get_py_cmd("areal.v2.inference_service.sglang.launch_server", args) @staticmethod def build_args( @@ -1999,7 +2104,7 @@ class AgentConfig: default_factory=lambda: ( SchedulingSpec( gpu=0, - cmd="python -m areal.experimental.agent_service.guard", + cmd="python -m areal.v2.agent_service.guard", ), ), metadata={ @@ -2056,6 +2161,30 @@ class AgentConfig: "choices": ["individual", "concat"], }, ) + message_preprocessors: list[str] = field( + default_factory=list, + metadata={ + "help": ( + "List of message preprocessor class paths applied, in order, to " + "Anthropic-compatible `/v1/messages` requests after translating " + "them to OpenAI-compatible requests. Native OpenAI " + "`/chat/completions` and `/responses` requests are not " + "preprocessed. Each entry is a dotted import path to a callable " + "class." + ), + }, + ) + prefix_matcher: str | None = field( + default=None, + metadata={ + "help": ( + "Dotted import path to a custom prefix matcher function for " + "InteractionCache parent-child matching. The function must accept " + "two list[dict] arguments (candidate prefix, full messages) and " + "return bool. When None, exact element-wise equality is used." + ), + }, + ) subproc_max_workers: int = field( default=4, metadata={ @@ -2842,6 +2971,14 @@ class BaseExperimentConfig: scheduler: SchedulerConfig = field(default_factory=SchedulerConfig) + post_exit_hook: str = field( + default="", + metadata={ + "help": "Shell command run after launcher shutdown. " + "LOG_DIR is injected; failures are logged and ignored." + }, + ) + def __post_init__(self): """Validate training configuration.""" if self.total_train_epochs <= 0: diff --git a/areal/api/reward_api.py b/areal/api/reward_api.py index 44bef1f692..5e05d49804 100644 --- a/areal/api/reward_api.py +++ b/areal/api/reward_api.py @@ -7,7 +7,9 @@ from collections.abc import Callable from concurrent.futures import ProcessPoolExecutor from concurrent.futures.process import BrokenProcessPool +from dataclasses import dataclass, field from functools import partial +from typing import Any from areal.utils import logging @@ -59,6 +61,44 @@ def reward_fn( """ +@dataclass +class RewardResult: + """Structured reward payload for workflows that expose process rewards. + + Parameters + ---------- + final_reward: + Sequence-level terminal reward. This is the compatibility anchor used + by existing workflows and PPO code paths. + step_rewards: + Optional reward value for each process step. + step_ends: + Optional 1-based completion-token end offsets for each step. Each entry + points to the completion token position where the corresponding + ``step_rewards`` value should be injected. + metadata: + Optional debug-only payload for logging or inspection. This field is + not consumed by the trainer in Stage 1. + """ + + final_reward: float + step_rewards: list[float] | None = None + step_ends: list[int] | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +def normalize_reward_result(value: float | RewardResult) -> RewardResult: + """Normalize legacy scalar rewards into ``RewardResult``. + + This keeps existing reward functions working unchanged while allowing new + implementations to return a structured process-reward payload. + """ + + if isinstance(value, RewardResult): + return value + return RewardResult(final_reward=float(value)) + + class AsyncRewardWrapper: """Wraps a synchronous reward function for async execution with timeout and retries. diff --git a/areal/engine/megatron_engine.py b/areal/engine/megatron_engine.py index 0713683dee..6e268451a8 100644 --- a/areal/engine/megatron_engine.py +++ b/areal/engine/megatron_engine.py @@ -5,9 +5,11 @@ import dataclasses import functools import gc +import json import math import os import re +import struct from collections.abc import Callable, Iterator from concurrent.futures import Future from contextlib import contextmanager @@ -558,6 +560,51 @@ def _build_hf_mcore_bridge(self): distribute_saved_activations=self.mcore_config.distribute_saved_activations, recompute_modules=self.mcore_config.recompute_modules, ) + + # Set MoE configuration overrides (aux-loss-free balancing, z-loss). + moe_extra_args: dict = { + "moe_token_dispatcher_type": self.mcore_config.moe_token_dispatcher_type, + "moe_permute_fusion": self.mcore_config.moe_permute_fusion, + "moe_router_fusion": self.mcore_config.moe_router_fusion, + "moe_shared_expert_overlap": self.mcore_config.moe_shared_expert_overlap, + "moe_router_bias_update_rate": self.mcore_config.moe_router_bias_update_rate, + } + if self.mcore_config.moe_router_dtype is not None: + moe_extra_args["moe_router_dtype"] = self.mcore_config.moe_router_dtype + if self.mcore_config.moe_z_loss_coeff is not None: + moe_extra_args["moe_z_loss_coeff"] = self.mcore_config.moe_z_loss_coeff + if self.mcore_config.moe_enable_deepep: + moe_extra_args["moe_enable_deepep"] = True + # Filter out args not accepted by the target TransformerConfig class. + accepted = { + f.name for f in dataclasses.fields(self.bridge.TransformerConfigClass) + } + moe_extra_args = {k: v for k, v in moe_extra_args.items() if k in accepted} + self.bridge.set_extra_args(**moe_extra_args) + + # Set precision and loss configuration (may not be supported by all + # model configs, e.g. MLATransformerConfig rejects enable_fp32_lm_head). + precision_args = {} + if self.mcore_config.enable_fp32_lm_head: + precision_args["enable_fp32_lm_head"] = True + if self.mcore_config.cross_entropy_loss_fusion: + precision_args["cross_entropy_loss_fusion"] = True + if precision_args: + skipped_precision_args = [ + k for k in precision_args if k not in accepted + ] + precision_args = { + k: v for k, v in precision_args.items() if k in accepted + } + if skipped_precision_args: + self.logger.warning( + "Some precision/loss args are not supported by this model " + f"config ({self.bridge.TransformerConfigClass.__name__}); " + f"skipping: {skipped_precision_args}" + ) + if precision_args: + self.bridge.set_extra_args(**precision_args) + self.logger.info( "Using mbridge to create models and hf model save/load in MegatronEngine." ) @@ -839,6 +886,7 @@ def forward_backward_batch( [torch.Tensor, dict[str, Any]], torch.Tensor | None ], forward_only: bool = False, + gather_cp_output: bool = False, ) -> None: self._ensure_ready() @@ -870,7 +918,13 @@ def forward_step(batch_iter, model): tree_attn_keys = list(tree_kwargs.keys()) cp_size = mpu.get_context_parallel_world_size() - cp_local = cp_size > 1 + # forward_batch (compute_logp / compute_values) passes + # gather_cp_output=True so CP-local outputs are gathered back to the + # full sequence length inside forward. This matches downstream + # labels / output_seqlens and fixes split_with_sizes mismatches when + # compute_logp runs with CP > 1. Train/eval keeps the default False + # value, so the CP-local loss path (_cp_local_labels) is unchanged. + cp_local = cp_size > 1 and not gather_cp_output output = packed_context_parallel_forward( model, @@ -1075,7 +1129,9 @@ def process_output(output: torch.Tensor, inputs: dict[str, Any]) -> None: outputs.append(result) return None - self.forward_backward_batch(mb_list, process_output, forward_only=True) + self.forward_backward_batch( + mb_list, process_output, forward_only=True, gather_cp_output=True + ) # Step 4: Aggregate, reorder, and broadcast outputs res = None @@ -1122,6 +1178,17 @@ def offload(self) -> None: ) self.get_device_stats().log("before offload model") + + # Discard gradient buffers via Megatron's native API *before* TMS pause. + # `DDP.offload_grad_buffers()` releases grad storage in-place + # (storage().resize_(0)); views like param.main_grad recover automatically + # on restore. Since grads are recomputed every step, they need no CPU + # backup. Doing this before pause() means the grad region has no physical + # memory left for TMS to back up. + if self.mcore_config.disable_grad_buffers_cpu_backup: + for m in self.model: + m.offload_grad_buffers(synchronize=False, empty_cache=False) + current_platform.clear_memory() torch_memory_saver.pause() @@ -1139,6 +1206,13 @@ def onload(self) -> None: """ torch_memory_saver.resume() + + # Reallocate gradient buffers released in offload(). resize_() restores + # storage and zeroes it; param.main_grad views become valid again. + if self.mcore_config.disable_grad_buffers_cpu_backup: + for m in self.model: + m.restore_grad_buffers(synchronize=False) + current_platform.clear_memory() # TODO: NCCL onload @@ -1954,10 +2028,17 @@ def _save_model_to_hf( base_model_name_or_path=base_model_path or self.config.path, ) else: + # When the MTP head was dropped (enable_mtp=False), the export + # yields no mtp.* tensors and strict=True would silently skip + # every source shard containing an MTP key -- discarding the + # non-MTP weights packed in those shards (e.g. lm_head). + # strict=False writes such shards with all present keys, so the + # export loses only the intentionally dropped MTP weights. self.bridge.save_hf_pretrained( self.model, path, source_path=base_model_path, + strict=not self._mtp_head_dropped, ) else: if self.mcore_config.use_mbridge_save: @@ -1989,10 +2070,96 @@ def _save_model_to_hf( tokenizer.save_pretrained(path) if processor is not None: processor.save_pretrained(path) + if self._mtp_head_dropped: + self._scrub_mtp_from_saved_config(path) + self._rebuild_index_from_saved_shards(path) current_platform.synchronize() dist.barrier(group=self.cpu_group) + @property + def _mtp_head_dropped(self) -> bool: + """True when the model declares an MTP head but enable_mtp left it unbuilt.""" + if self.bridge_cls != "megatron-bridge" or self.mcore_config.enable_mtp: + return False + text_config = getattr(self.hf_config, "text_config", self.hf_config) + return any( + bool(getattr(text_config, key, 0)) + for key in ("mtp_num_hidden_layers", "num_nextn_predict_layers") + ) + + def _scrub_mtp_from_saved_config(self, path: str) -> None: + """Zero MTP layer counts in an exported config.json so it matches the + MTP-stripped weights (the head is dropped when enable_mtp=False).""" + cfg_path = os.path.join(path, "config.json") + if not os.path.exists(cfg_path): + return + with open(cfg_path) as f: + cfg = json.load(f) + + def _walk(node: Any) -> bool: + changed = False + if isinstance(node, dict): + for key, value in node.items(): + if ( + key in ("mtp_num_hidden_layers", "num_nextn_predict_layers") + and value + ): + node[key] = 0 + changed = True + else: + changed |= _walk(value) + elif isinstance(node, list): + for item in node: + changed |= _walk(item) + return changed + + if _walk(cfg): + with open(cfg_path, "w") as f: + json.dump(cfg, f, indent=2, sort_keys=True) + self.logger.info( + "Exported checkpoint is MTP-stripped (enable_mtp=False); " + "zeroed MTP layer counts in %s.", + cfg_path, + ) + + def _rebuild_index_from_saved_shards(self, path: str) -> None: + """Rewrite model.safetensors.index.json from the shard files' actual + contents. megatron-bridge's strict=False save marks each shard's full + expected key set as saved, leaving ghost entries for the dropped mtp.* + tensors and a stale metadata.total_size.""" + index_path = os.path.join(path, "model.safetensors.index.json") + if not os.path.exists(index_path): + return + weight_map: dict[str, str] = {} + total_size = 0 + for filename in sorted(os.listdir(path)): + if not filename.endswith(".safetensors"): + continue + with open(os.path.join(path, filename), "rb") as f: + header_len = struct.unpack(" None: assert self.model is not None, "Model is not initialized." @@ -2159,6 +2326,10 @@ def _compute_logprobs_and_loss( vocab_min_logits, vocab_max_logits = gather_packed_tree_vocab_stats( output, trie_node ) + # Tree training only supports packed min/max vocab stats; mean/norm + # would need per-sequence unpacking, so leave them unset. + vocab_mean_logits = None + vocab_norm_logits = None logprobs, entropy = gather_packed_tree_logprobs_entropy( output, trie_node, @@ -2185,6 +2356,8 @@ def _compute_logprobs_and_loss( ) vocab_min_logits = output.detach().min(-1).values.float() vocab_max_logits = output.detach().max(-1).values.float() + vocab_mean_logits = output.detach().float().mean(-1) + vocab_norm_logits = output.detach().float().norm(dim=-1) if cp_padded_cu_seqlens is not None: logprobs = reassemble_cp_packed_logprobs( logprobs, cp_padded_cu_seqlens @@ -2198,6 +2371,12 @@ def _compute_logprobs_and_loss( vocab_max_logits = reassemble_cp_packed_logprobs( vocab_max_logits, cp_padded_cu_seqlens ) + vocab_mean_logits = reassemble_cp_packed_logprobs( + vocab_mean_logits, cp_padded_cu_seqlens + ) + vocab_norm_logits = reassemble_cp_packed_logprobs( + vocab_norm_logits, cp_padded_cu_seqlens + ) cp_padding_length = inputs.get("_cp_padding_length", 0) cp_old_cu_seqlens = inputs.get("_cp_old_cu_seqlens") logprobs = unpad_logits( @@ -2224,6 +2403,18 @@ def _compute_logprobs_and_loss( cp_padded_cu_seqlens, cp_old_cu_seqlens, ) + vocab_mean_logits = unpad_logits( + vocab_mean_logits, + cp_padding_length, + cp_padded_cu_seqlens, + cp_old_cu_seqlens, + ) + vocab_norm_logits = unpad_logits( + vocab_norm_logits, + cp_padding_length, + cp_padded_cu_seqlens, + cp_old_cu_seqlens, + ) inputs = { k: v for k, v in inputs.items() if not k.startswith("_cp_") } @@ -2233,6 +2424,8 @@ def _compute_logprobs_and_loss( inputs, vocab_min_logits=vocab_min_logits, vocab_max_logits=vocab_max_logits, + vocab_mean_logits=vocab_mean_logits, + vocab_norm_logits=vocab_norm_logits, ) else: values = output.squeeze(-1) diff --git a/areal/engine/megatron_utils/checkpointer.py b/areal/engine/megatron_utils/checkpointer.py index 580e2bd6a2..6bdc55e2e0 100644 --- a/areal/engine/megatron_utils/checkpointer.py +++ b/areal/engine/megatron_utils/checkpointer.py @@ -278,6 +278,7 @@ def generate_state_dict( with_model: bool = True, with_optimizer: bool = True, with_rng: bool = True, + is_loading: bool = False, ): # For save dist checkpointing state_dict = {} @@ -298,7 +299,24 @@ def generate_state_dict( # Optimizer State Dict if with_optimizer: torch.distributed.barrier() - optimizer_sharded_states = self.optimizer.sharded_state_dict(state_dict) + # megatron-core v0.14+ removed flattened_range support (Megatron-LM + # PR #2126), but the sharded_state_dict default + # (fully_sharded_model_space) still emits it, so saving optimizer + # state fails on the pinned 0.17.0. dp_reshardable is upstream's + # current default. Trade-off: the optimizer state (not the model + # weights) becomes reshardable only along DP -- load hard-asserts + # the same bucket layout (per_bucket_numel_unpadded), so save and + # load must use identical TP/PP. Fine today since recover enforces + # the same topology; switch to fully_reshardable if cross-topology + # resume is ever needed (also flattened_range-free, at the cost of + # gathering optimizer state). is_loading=True pre-allocates + # exp_avg/exp_avg_sq so the load template requests them -- + # otherwise DCP silently drops the moments on resume. + optimizer_sharded_states = self.optimizer.sharded_state_dict( + state_dict, + is_loading=is_loading, + metadata={"distrib_optim_sharding_type": "dp_reshardable"}, + ) state_dict["optimizer"] = optimizer_sharded_states if self.lr_scheduler is not None: @@ -353,7 +371,7 @@ def load_checkpoint( # Get State Dict for loading sharded_state_dict = self.generate_state_dict( - with_model, with_optimizer, with_rng + with_model, with_optimizer, with_rng, is_loading=True ) # Load Dist Checkpointing state_dict = load_dist_checkpointing( diff --git a/areal/engine/sglang_remote.py b/areal/engine/sglang_remote.py index a2b1677aa3..7221326b76 100644 --- a/areal/engine/sglang_remote.py +++ b/areal/engine/sglang_remote.py @@ -587,7 +587,7 @@ def export_stats(self) -> dict[str, float]: @classmethod def as_controller(cls, config: InferenceEngineConfig, scheduler: Scheduler): if config._version == "v2": - from areal.experimental.inference_service.controller.controller import ( + from areal.v2.inference_service.controller.controller import ( RolloutControllerV2, ) diff --git a/areal/engine/vllm_remote.py b/areal/engine/vllm_remote.py index c0c4a4a59e..42aee39dc6 100644 --- a/areal/engine/vllm_remote.py +++ b/areal/engine/vllm_remote.py @@ -58,11 +58,14 @@ def build_generation_request( "stop_token_ids": stop_token_ids, "ignore_eos": gconfig.ignore_eos, "skip_special_tokens": gconfig.skip_special_tokens, + "frequency_penalty": gconfig.frequency_penalty, "return_tokens_as_token_ids": True, "logprobs": 0, "use_beam_search": gconfig.use_beam_search, "stream": False, } + if gconfig.stop: + payload["stop"] = gconfig.stop if with_lora: lora_name = gconfig.lora_name @@ -539,7 +542,7 @@ def export_stats(self) -> dict[str, float]: @classmethod def as_controller(cls, config: InferenceEngineConfig, scheduler: Scheduler): if config._version == "v2": - from areal.experimental.inference_service.controller.controller import ( + from areal.v2.inference_service.controller.controller import ( RolloutControllerV2, ) diff --git a/areal/experimental/agent_service/data_proxy/app.py b/areal/experimental/agent_service/data_proxy/app.py deleted file mode 100644 index fbf51646ba..0000000000 --- a/areal/experimental/agent_service/data_proxy/app.py +++ /dev/null @@ -1,155 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -"""Data Proxy — stateful session proxy between Gateway and Worker.""" - -from __future__ import annotations - -import asyncio -import time -from dataclasses import dataclass, field -from typing import Any - -import httpx -from fastapi import FastAPI - -from areal.utils import logging - -from .config import DataProxyConfig - -logger = logging.getLogger("AgentDataProxy") - - -@dataclass -class _SessionData: - history: list[dict[str, Any]] = field(default_factory=list) - metadata: dict[str, Any] = field(default_factory=dict) - last_active: float = field(default_factory=time.monotonic) - - -def create_data_proxy_app(config: DataProxyConfig) -> FastAPI: - app = FastAPI(title="AReaL Data Proxy") - sessions: dict[str, _SessionData] = {} - http_client = httpx.AsyncClient(timeout=config.request_timeout) - - async def _close_worker_session(session_key: str) -> None: - try: - await http_client.post( - f"{config.worker_addr}/session/{session_key}/close", timeout=5 - ) - except Exception: - logger.debug("Failed to close worker session %s", session_key) - - async def _reap_idle_sessions() -> None: - while True: - await asyncio.sleep(60) - now = time.monotonic() - stale = [ - k - for k, s in sessions.items() - if now - s.last_active > config.session_timeout - ] - for k in stale: - del sessions[k] - await _close_worker_session(k) - if stale: - logger.info("Reaped %d idle sessions", len(stale)) - - @app.on_event("startup") - async def startup(): - app.state.reaper_task = asyncio.create_task(_reap_idle_sessions()) - - @app.on_event("shutdown") - async def shutdown(): - await http_client.aclose() - - @app.get("/health") - async def health(): - return { - "status": "ok", - "active_sessions": len(sessions), - "worker_addr": config.worker_addr, - } - - @app.post("/session/{session_key}/turn") - async def turn(session_key: str, body: dict[str, Any]): - session = sessions.get(session_key) - if session is None: - session = _SessionData() - sessions[session_key] = session - - message = body.get("message", "") - run_id = body.get("run_id", "") - queue_mode = body.get("queue_mode", "collect") - metadata = body.get("metadata", {}) - - worker_request = { - "message": message, - "session_key": session_key, - "run_id": run_id, - "history": session.history.copy(), - "queue_mode": queue_mode, - "metadata": metadata, - } - - resp = await http_client.post(f"{config.worker_addr}/run", json=worker_request) - resp.raise_for_status() - result = resp.json() - - session.history.append({"role": "user", "content": message}) - - call_counter = 0 - for evt in result.get("events", []): - if evt.get("type") == "tool_call": - call_id = f"call_{evt.get('name', '')}_{run_id}_{call_counter}" - call_counter += 1 - session.history.append( - { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": call_id, - "type": "function", - "function": { - "name": evt.get("name", ""), - "arguments": evt.get("args", ""), - }, - } - ], - } - ) - elif evt.get("type") == "tool_result": - result_call_id = ( - f"call_{evt.get('name', '')}_{run_id}_{call_counter - 1}" - if call_counter > 0 - else f"call_{evt.get('name', '')}_{run_id}_0" - ) - session.history.append( - { - "role": "tool", - "tool_call_id": result_call_id, - "content": evt.get("result", ""), - } - ) - - summary = result.get("summary", "") - if summary: - session.history.append({"role": "assistant", "content": summary}) - - session.last_active = time.monotonic() - return result - - @app.post("/session/{session_key}/close") - async def close_session(session_key: str): - sessions.pop(session_key, None) - await _close_worker_session(session_key) - return {"status": "ok"} - - @app.get("/session/{session_key}/history") - async def get_history(session_key: str): - session = sessions.get(session_key) - if session is None: - return {"history": []} - return {"history": session.history} - - return app diff --git a/areal/experimental/agent_service/gateway/bridge.py b/areal/experimental/agent_service/gateway/bridge.py deleted file mode 100644 index e47c8d8f35..0000000000 --- a/areal/experimental/agent_service/gateway/bridge.py +++ /dev/null @@ -1,172 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -"""OpenResponses HTTP bridge — translates POST /v1/responses to DataProxy turns.""" - -from __future__ import annotations - -import uuid -from abc import ABC, abstractmethod -from typing import Any - -import httpx -from fastapi import Depends, FastAPI, Request -from fastapi.responses import JSONResponse - -from areal.utils import logging - -from ..auth import DEFAULT_ADMIN_API_KEY, admin_headers, make_admin_dependency -from ..protocol import generate_run_id - -logger = logging.getLogger("AgentBridge") - - -class AgentBridge(ABC): - @abstractmethod - async def handle_request(self, request: Request) -> Any: ... - - -class OpenResponsesBridge(AgentBridge): - def __init__( - self, router_addr: str, admin_api_key: str = DEFAULT_ADMIN_API_KEY - ) -> None: - self._router_addr = router_addr - self._auth_headers = admin_headers(admin_api_key) - self._http = httpx.AsyncClient(timeout=600.0) - - async def close(self) -> None: - await self._http.aclose() - - async def handle_request(self, request: Request) -> Any: - body = await request.json() - - input_items: list[dict[str, Any]] = body.get("input", []) - instructions: str = body.get("instructions", "") - model: str = body.get("model", "") - user: str = body.get("user", "") - - if not user: - return JSONResponse( - { - "error": { - "message": "'user' field is required for session affinity", - "type": "invalid_request", - } - }, - status_code=400, - ) - - message = self._extract_message(input_items, instructions) - session_key = self._derive_session_key(user, model) - run_id = generate_run_id() - response_id = f"resp-{uuid.uuid4().hex[:12]}" - - metadata = { - "input": input_items, - "instructions": instructions, - "tools": body.get("tools", []), - "model": model, - "idempotencyKey": response_id, - **body.get("metadata", {}), - } - - try: - route_resp = await self._http.post( - f"{self._router_addr}/route", - json={"session_key": session_key}, - headers=self._auth_headers, - ) - route_resp.raise_for_status() - data_proxy_addr = route_resp.json()["data_proxy_addr"] - - turn_resp = await self._http.post( - f"{data_proxy_addr}/session/{session_key}/turn", - json={ - "message": message, - "run_id": run_id, - "queue_mode": "collect", - "metadata": metadata, - }, - ) - turn_resp.raise_for_status() - result = turn_resp.json() - - output_items: list[dict[str, Any]] = [] - summary = result.get("summary", "") - if summary: - output_items.append( - { - "type": "message", - "role": "assistant", - "content": [{"type": "output_text", "text": summary}], - } - ) - - for evt in result.get("events", []): - if evt.get("type") == "tool_call": - output_items.append( - { - "type": "function_call", - "name": evt.get("name", ""), - "arguments": evt.get("args", ""), - } - ) - - return JSONResponse( - { - "id": response_id, - "object": "response", - "status": "completed", - "output": output_items, - "model": model, - "metadata": result.get("metadata", {}), - } - ) - except Exception as exc: - logger.error("OpenResponses request failed: %s", exc) - return JSONResponse( - {"error": {"message": str(exc), "type": "server_error"}}, - status_code=500, - ) - - @staticmethod - def _extract_message(input_items: list[dict[str, Any]], instructions: str) -> str: - parts: list[str] = [] - if instructions: - parts.append(instructions) - for item in input_items: - if item.get("type") == "message": - content = item.get("content", "") - if isinstance(content, list): - for block in content: - if ( - isinstance(block, dict) - and block.get("type") == "input_text" - ): - parts.append(block.get("text", "")) - elif isinstance(content, str): - parts.append(content) - elif item.get("type") == "function_call_output": - parts.append(f"[tool result] {item.get('output', '')}") - return "\n".join(parts) - - @staticmethod - def _derive_session_key(user: str, model: str) -> str: - if user: - return f"agent:{model or 'default'}:{user}" - return f"agent:{model or 'default'}:{uuid.uuid4().hex[:8]}" - - -def mount_bridge( - app: FastAPI, - bridge: OpenResponsesBridge, - admin_api_key: str = DEFAULT_ADMIN_API_KEY, -) -> None: - auth = make_admin_dependency(admin_api_key) - - @app.post("/v1/responses", dependencies=[Depends(auth)]) - async def responses_endpoint(request: Request): - return await bridge.handle_request(request) - - @app.on_event("shutdown") - async def shutdown_bridge(): - await bridge.close() diff --git a/areal/experimental/agent_service/types.py b/areal/experimental/agent_service/types.py deleted file mode 100644 index 7290634e67..0000000000 --- a/areal/experimental/agent_service/types.py +++ /dev/null @@ -1,66 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -"""Public types for the Agent Service protocol.""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from typing import Any, Protocol, runtime_checkable - -from .protocol import QueueMode - - -@dataclass -class AgentRequest: - """Structured request passed to the agent. - - Core fields are stable protocol-level attributes. Framework-specific - parameters should go in *metadata*. - """ - - message: str - session_key: str - run_id: str - history: list[dict[str, Any]] = field(default_factory=list) - queue_mode: QueueMode = QueueMode.COLLECT - metadata: dict[str, Any] = field(default_factory=dict) - - -@dataclass -class AgentResponse: - """Structured result returned by the agent.""" - - summary: str = "" - metadata: dict[str, Any] = field(default_factory=dict) - - -class EventEmitter(Protocol): - """Callback interface for streaming events from agent to caller.""" - - async def emit_delta(self, text: str) -> None: ... - async def emit_tool_call(self, name: str, args: str) -> None: ... - async def emit_tool_result(self, name: str, result: str) -> None: ... - - -@runtime_checkable -class AgentRunnable(Protocol): - """Minimal protocol for pluggable agent implementations. - - Agent classes are loaded via - :func:`~areal.utils.dynamic_import.import_from_string` at worker startup. - The framework handles its own tool execution, memory, and LLM - interaction — the Agent Service only provides session lifecycle and - event streaming. - - Reward computation is **not** part of this interface. Rewards are - calculated externally by the training pipeline (via reward functions - applied to exported trajectories), following AReaL's standard RLVR - pattern. - """ - - async def run( - self, - request: AgentRequest, - *, - emitter: EventEmitter, - ) -> AgentResponse: ... diff --git a/areal/experimental/openai/cache.py b/areal/experimental/openai/cache.py index 20fee73ca8..33888a4d6c 100644 --- a/areal/experimental/openai/cache.py +++ b/areal/experimental/openai/cache.py @@ -1,17 +1,55 @@ # SPDX-License-Identifier: Apache-2.0 +import json +import os import threading +import time from collections import OrderedDict -from typing import Any +from collections.abc import Callable +from pathlib import Path +from typing import Any, Protocol, runtime_checkable from areal.experimental.openai.types import InteractionWithTokenLogpReward from areal.utils import logging logger = logging.getLogger("OpenAICache") +# Debug-only: dumping mismatched parent/child messages is OFF by default because +# the payloads can contain full conversations. Set the dump dir env var to opt in. +_MISMATCH_DUMP_DIR_ENV = "AREAL_OPENAI_CACHE_MISMATCH_DUMP_DIR" +_MISMATCH_DUMP_LIMIT_ENV = "AREAL_OPENAI_CACHE_MISMATCH_DUMP_LIMIT" +_DEFAULT_MISMATCH_DUMP_LIMIT = 20 + + +@runtime_checkable +class PrefixMatcher(Protocol): + """Protocol for custom message prefix matching functions.""" + + def __call__(self, a: list[dict], b: list[dict]) -> bool: ... + + +def default_prefix_matcher(a: list[dict], b: list[dict]) -> bool: + """Default exact prefix matcher: ``b[:len(a)] == a``.""" + if len(a) > len(b): + return False + return b[: len(a)] == a + class InteractionCache(OrderedDict[str, InteractionWithTokenLogpReward]): - def __init__(self, *args, **kwargs): + def __init__( + self, + *args, + session_id: str = "unknown", + prefix_matcher: PrefixMatcher + | Callable[[list[dict], list[dict]], bool] + | None = None, + **kwargs, + ): + self._match_fail_count = 0 + self._session_id = session_id + self._prefix_matcher: Callable[[list[dict], list[dict]], bool] = ( + prefix_matcher if prefix_matcher is not None else default_prefix_matcher + ) super().__init__(*args, **kwargs) self._apply_reward_discount_called = False self._total_reward = 0.0 @@ -29,7 +67,10 @@ def __deepcopy__(self, memo): assert len(self) == 0, ( f"InteractionCache must be empty when deep-copied, but has {len(self)} items" ) - new = InteractionCache() + new = InteractionCache( + session_id=self._session_id, + prefix_matcher=self._prefix_matcher, + ) memo[id(self)] = new return new @@ -115,12 +156,6 @@ def __setitem__( "Interaction messages must be set to find parent relationship." ) - def _is_prefix(a: list[dict], b: list[dict]) -> bool: - # True if a is a prefix of b - if len(a) > len(b): - return False - return b[: len(a)] == a - def _is_similar_on_last_message( a: list[dict], b: list[dict] ) -> tuple[bool, dict[str, Any] | None, dict[str, Any] | None]: @@ -143,13 +178,12 @@ def _is_similar_on_last_message( } return True, diff_a_message, diff_b_message - # Construct parent-child relationships using longest prefix rule - # Sort potential parents by message length to find the longest prefix match first. + # Construct parent-child relationships using longest prefix rule. interactions = sorted( self.values(), key=lambda x: len(x.messages), reverse=True ) + child_msgs = value.messages - # Find parent for the new interaction for parent in interactions: # Skip interactions that are still being processed (output_message_list not set yet) # This can happen with concurrent requests where a streaming request hasn't @@ -157,24 +191,103 @@ def _is_similar_on_last_message( if parent.output_message_list is None or parent.messages is None: continue parent_data = parent.messages + parent.output_message_list - if _is_prefix(parent_data, value.messages): + if self._prefix_matcher(parent_data, child_msgs): value.parent = parent break - elif _is_prefix(parent.messages, value.messages): + elif self._prefix_matcher(parent.messages, child_msgs): + self._match_fail_count += 1 + + logger.warning( + "Prefix mismatch (occurrence %d, session=%s): " + "parent.messages (len=%d) is a prefix of child (len=%d), " + "but parent_data (messages+output, len=%d) is not.", + self._match_fail_count, + self._session_id, + len(parent.messages), + len(child_msgs), + len(parent_data), + ) + self._dump_mismatch( + parent_data=parent_data, + child_msgs=child_msgs, + parent_id=parent.interaction_id, + child_key=key, + ) + is_similar, diff_a, diff_b = _is_similar_on_last_message( - parent_data, value.messages + parent_data, child_msgs ) if is_similar: logger.warning( "Found a parent interaction with similar last message content, " - "but not a strict prefix match. If you wish to use concat mode and build a conversation tree:\n" + "but not a strict prefix match (occurrence %d, " + "session=%s, parent_len=%d, child_len=%d). " + "If you wish to use concat mode and build a conversation tree:\n" "1. For completion, append `chat_completion.choices[0].message.model_dump()` to your messages.\n" "2. For response, extend `[o.model_dump() for o in response.output]` to your messages.\n" - f"Different keys in parent last message: {diff_a}\n" - f"Different keys in child last message: {diff_b}\n" + "Different keys in parent last message: %s\n" + "Different keys in child last message: %s", + self._match_fail_count, + self._session_id, + len(parent_data), + len(child_msgs), + diff_a, + diff_b, ) super().__setitem__(key, value) + def _dump_mismatch( + self, + parent_data: list[dict], + child_msgs: list[dict], + parent_id: str | None, + child_key: str, + ) -> None: + """Dump mismatched parent/child messages to a JSON file for debugging.""" + dump_dir_env = os.environ.get(_MISMATCH_DUMP_DIR_ENV) + if not dump_dir_env: + return + try: + dump_limit = int( + os.environ.get( + _MISMATCH_DUMP_LIMIT_ENV, str(_DEFAULT_MISMATCH_DUMP_LIMIT) + ) + ) + if dump_limit > 0 and self._match_fail_count > dump_limit: + return + + dump_dir = Path(dump_dir_env) + dump_dir.mkdir(parents=True, exist_ok=True) + + ts = time.strftime("%Y%m%d_%H%M%S") + filename = f"mismatch_{self._session_id}_{self._match_fail_count}_{ts}.json" + dump_path = dump_dir / filename + + first_diff_idx = None + for i, (pm, cm) in enumerate(zip(parent_data, child_msgs)): + if pm != cm: + first_diff_idx = i + break + + dump = { + "session_id": self._session_id, + "occurrence": self._match_fail_count, + "parent_id": parent_id, + "child_key": child_key, + "parent_len": len(parent_data), + "child_len": len(child_msgs), + "first_diff_idx": first_diff_idx, + "parent_data": parent_data, + "child_msgs": child_msgs, + } + + with open(dump_path, "w", encoding="utf-8") as f: + json.dump(dump, f, indent=2, ensure_ascii=False, default=str) + + logger.info("Mismatch dump saved to %s", dump_path) + except Exception as e: + logger.warning("Failed to dump mismatch: %s", e) + def export_interactions( self, style: str, reward_discount: float | None = None ) -> dict[str, InteractionWithTokenLogpReward]: diff --git a/areal/experimental/openai/client.py b/areal/experimental/openai/client.py index 016bd738c5..dc68117fed 100644 --- a/areal/experimental/openai/client.py +++ b/areal/experimental/openai/client.py @@ -76,6 +76,8 @@ async def agenerate(self, req: ModelRequest) -> ModelResponse: logger = logging.getLogger("OpenAIClient") +_DEFAULT_MAX_TOTAL_TOKENS = 32768 + def _ensure_message_dict_list( name: str, @@ -343,6 +345,55 @@ def _build_messages_list(item: dict) -> list[dict]: return messages_list +def _resolve_max_total_tokens( + prompt_len: int, + max_new_tokens: int, + engine_max_tokens: int | None, +) -> int: + # Fall back to a fixed context-length ceiling when the deployment does not + # configure engine_max_tokens, so a large client max_tokens cannot push + # prompt + generation past the backend model's context window. + cap = engine_max_tokens or _DEFAULT_MAX_TOTAL_TOKENS + return min(prompt_len + max_new_tokens, cap) + + +def _parse_tool_call_arguments(messages: list[dict]) -> list[dict]: + """Return a new message list with tool_call arguments parsed from JSON strings + to dicts. Some chat templates (e.g. GLM-5.1) iterate over arguments with + .items(), which fails when arguments is a JSON string per OpenAI API convention. + + Only creates new dicts where modifications are needed; unaffected messages + are shared with the input list to avoid expensive deep copies. + """ + result = [] + for msg in messages: + tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else None + if not tool_calls: + result.append(msg) + continue + new_tool_calls = [] + modified = False + for tc in tool_calls: + fn = tc.get("function") if isinstance(tc, dict) else None + if fn is None: + new_tool_calls.append(tc) + continue + args = fn.get("arguments") + if isinstance(args, str): + try: + parsed = json.loads(args) + tc = {**tc, "function": {**fn, "arguments": parsed}} + modified = True + except (json.JSONDecodeError, TypeError): + pass + new_tool_calls.append(tc) + if modified: + result.append({**msg, "tool_calls": new_tool_calls}) + else: + result.append(msg) + return result + + def concat_prompt_token_ids_with_parent( message_list: list[dict], parent: InteractionWithTokenLogpReward | None, @@ -385,6 +436,7 @@ def concat_prompt_token_ids_with_parent( parent_tokens += [eos_token_id] all_message_list += message_list + all_message_list = _parse_tool_call_arguments(all_message_list) all_tokens = apply_chat_template( tokenizer, @@ -610,6 +662,7 @@ async def create( has_images = len(image_data) > 0 tokenizer_messages = messages_for_tokenizer if has_images else messages_list + tokenizer_messages = _parse_tool_call_arguments(tokenizer_messages) if self.chat_template_type == "hf": prompt_token_ids = apply_chat_template( self.tokenizer, @@ -720,6 +773,11 @@ async def create( n_samples=n, temperature=temp, max_new_tokens=max_new_tokens, + max_tokens=_resolve_max_total_tokens( + prompt_len=len(prompt_token_ids), + max_new_tokens=max_new_tokens, + engine_max_tokens=self.engine_max_tokens, + ), top_p=top_p_val, stop=stop_tokens, greedy=temp == 0, @@ -859,9 +917,17 @@ async def _create_stream( ) # Tool calls chunks (if any) + # Split each tool call into two chunks (name then arguments) to + # match standard OpenAI streaming behavior. LiteLLM's Anthropic + # streaming adapter treats the first chunk with a function name as + # a content_block_start trigger and discards the processed delta + # from that same chunk. If name and arguments are combined in a + # single chunk, the arguments are lost and Anthropic clients see + # input={}. if tool_calls: for idx, tool_call in enumerate(tool_calls): tool_call = cast(ChatCompletionMessageFunctionToolCall, tool_call) + # Chunk 1: name + id, with empty arguments. yield ChatCompletionChunk( id=completion_id, choices=[ @@ -874,6 +940,30 @@ async def _create_stream( type="function", function=ChoiceDeltaToolCallFunction( name=tool_call.function.name, + arguments="", + ), + ) + ] + ), + index=0, + finish_reason=None, + ) + ], + created=current_time, + model="None", + object="chat.completion.chunk", + ) + # Chunk 2: arguments only, emitted as input_json_delta by + # Anthropic adapters after the tool_use block has started. + yield ChatCompletionChunk( + id=completion_id, + choices=[ + ChunkChoice( + delta=ChoiceDelta( + tool_calls=[ + ChoiceDeltaToolCall( + index=idx, + function=ChoiceDeltaToolCallFunction( arguments=tool_call.function.arguments, ), ) @@ -1018,6 +1108,7 @@ async def create( has_images = len(image_data) > 0 tokenizer_messages = messages_for_tokenizer if has_images else messages_list + tokenizer_messages = _parse_tool_call_arguments(tokenizer_messages) if self.chat_template_type == "hf": prompt_token_ids = apply_chat_template( self.tokenizer, @@ -1081,6 +1172,11 @@ async def create( n_samples=1, temperature=temp, max_new_tokens=max_new_tokens, + max_tokens=_resolve_max_total_tokens( + prompt_len=len(prompt_token_ids), + max_new_tokens=max_new_tokens, + engine_max_tokens=self.engine_max_tokens, + ), top_p=top_p_val, stop=stop, greedy=temp == 0, diff --git a/areal/experimental/openai/proxy/proxy_rollout_server.py b/areal/experimental/openai/proxy/proxy_rollout_server.py index 38915590e4..5d5453d92d 100644 --- a/areal/experimental/openai/proxy/proxy_rollout_server.py +++ b/areal/experimental/openai/proxy/proxy_rollout_server.py @@ -110,6 +110,14 @@ def _warn_once(msg: str) -> None: _api_key_to_session: dict[str, str] = {} _session_to_api_key: dict[str, str] = {} # Reverse mapping for O(1) cleanup +# Pluggable message preprocessors loaded from config at setup time. +# Applied in order after Anthropic-to-OpenAI translation, before content +# reaches the ArealOpenAI client. +_message_preprocessors: list = [] + +# Pluggable prefix matcher for InteractionCache parent-child matching. +_prefix_matcher = None + # Server address (set at startup) _server_host: str = "0.0.0.0" _server_port: int = 8000 @@ -262,6 +270,7 @@ async def alloc_ports(raw_request: Request): def _setup_openai_client(): global _openai_client, _session_timeout_seconds, _admin_api_key + global _message_preprocessors, _prefix_matcher config = _engine.config tokenizer = load_hf_tokenizer(config.tokenizer_path) agent_cfg = config.agent @@ -292,6 +301,18 @@ def _setup_openai_client(): with _lock: _admin_api_key = agent_cfg.admin_api_key + _message_preprocessors = [] + for path in agent_cfg.message_preprocessors: + cls = import_from_string(path) + _message_preprocessors.append(cls()) + logger.info("Loaded message preprocessor: %s", path) + + if agent_cfg.prefix_matcher: + _prefix_matcher = import_from_string(agent_cfg.prefix_matcher) + logger.info("Loaded prefix matcher: %s", agent_cfg.prefix_matcher) + else: + _prefix_matcher = None + @app.post("/configure") async def configure(raw_request: Request): @@ -464,7 +485,10 @@ def start_session(request: StartSessionRequest) -> StartSessionResponse: session_api_key = secrets.token_urlsafe(32) _capacity -= 1 - _session_cache[session_id] = SessionData(session_id=session_id) + _session_cache[session_id] = SessionData( + session_id=session_id, + prefix_matcher=_prefix_matcher, + ) _api_key_to_session[session_api_key] = session_id _session_to_api_key[session_id] = session_api_key @@ -698,6 +722,19 @@ async def responses( ) +def _flatten_content_lists(messages: list[dict]) -> None: + """Flatten Anthropic content block lists to strings in-place.""" + for msg in messages: + if isinstance(msg.get("content"), list): + text_parts = [] + for block in msg["content"]: + if isinstance(block, dict) and block.get("type") == "text": + text_parts.append(block.get("text", "")) + elif isinstance(block, str): + text_parts.append(block) + msg["content"] = "\n".join(text_parts) + + def _translate_anthropic_to_openai_request(anthropic_request: dict[str, Any]) -> dict: """Translate an Anthropic Messages API request to OpenAI format.""" openai_request = _adapter.translate_completion_input_params( @@ -707,20 +744,10 @@ def _translate_anthropic_to_openai_request(anthropic_request: dict[str, Any]) -> raise ValueError("Failed to translate request") openai_request = dict(openai_request) - # Fix message content if it's a list (Anthropic format with content blocks) - # LiteLLM's adapter may not properly convert content from list to string - # Claude Code CLI sends content as: [{"type":"text","text":"...","cache_control":{...}}, ...] if "messages" in openai_request: - for msg in openai_request["messages"]: - if isinstance(msg.get("content"), list): - # Convert list of content blocks to string - text_parts = [] - for block in msg["content"]: - if isinstance(block, dict) and block.get("type") == "text": - text_parts.append(block.get("text", "")) - elif isinstance(block, str): - text_parts.append(block) - msg["content"] = "\n".join(text_parts) + _flatten_content_lists(openai_request["messages"]) + for preprocessor in _message_preprocessors: + openai_request["messages"] = preprocessor(openai_request["messages"]) return openai_request diff --git a/areal/experimental/openai/proxy/server.py b/areal/experimental/openai/proxy/server.py index 6e160c4644..54e4885e21 100644 --- a/areal/experimental/openai/proxy/server.py +++ b/areal/experimental/openai/proxy/server.py @@ -66,11 +66,14 @@ class ExportTrajectoriesResponse(BaseModel): class SessionData: """Data associated with a single RL session.""" - def __init__(self, session_id: str): + def __init__(self, session_id: str, prefix_matcher=None): self.session_id = session_id self._completed = False - self._completions = InteractionCache() + self._completions = InteractionCache( + session_id=session_id, + prefix_matcher=prefix_matcher, + ) self._completed_event = threading.Event() self._start_time = time.time() self._last_access_time = time.time() diff --git a/areal/experimental/openai/tool_call_parser.py b/areal/experimental/openai/tool_call_parser.py index 850b020bd0..4425e241a5 100644 --- a/areal/experimental/openai/tool_call_parser.py +++ b/areal/experimental/openai/tool_call_parser.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 +import json +import re import traceback import uuid from types import SimpleNamespace @@ -15,6 +17,19 @@ logger = logging.getLogger("ToolCallParser") +_QWEN3_CODER_TOOL_CALL_RE = re.compile( + r"\s*(?P.*?)\s*", + re.DOTALL, +) +_QWEN3_CODER_FUNCTION_RE = re.compile( + r"[^>\s]+)>\s*(?P.*?)\s*", + re.DOTALL, +) +_QWEN3_CODER_PARAMETER_RE = re.compile( + r"[^>\s]+)>(?P.*?)", + re.DOTALL, +) + _SGLANG_TO_VLLM_TOOL_PARSER: dict[str, str] = { "qwen": "qwen3_xml", "qwen25": "qwen3_xml", @@ -58,6 +73,187 @@ def _detect_think_and_return_ori_think( return think_start_token + reasoning_text + think_end_token, normal_text +def _iter_tool_definitions(tools: list[Any]) -> list[dict[str, Any]]: + """Return OpenAI chat/responses tool definitions in a common shape.""" + tool_defs: list[dict[str, Any]] = [] + for tool in tools: + if not isinstance(tool, dict): + continue + if isinstance(tool.get("function"), dict): + tool_defs.append(tool["function"]) + elif tool.get("type") == "function": + tool_defs.append(tool) + return tool_defs + + +def _tool_argument_schemas(tools: list[Any]) -> dict[str, dict[str, dict[str, Any]]]: + schemas: dict[str, dict[str, dict[str, Any]]] = {} + for tool_def in _iter_tool_definitions(tools): + name = tool_def.get("name") + parameters = tool_def.get("parameters") + if not isinstance(name, str) or not isinstance(parameters, dict): + continue + properties = parameters.get("properties") + if isinstance(properties, dict): + schemas[name] = { + key: value + for key, value in properties.items() + if isinstance(value, dict) + } + return schemas + + +def _clean_qwen3_coder_parameter(raw_value: str) -> str: + if raw_value.startswith("\n"): + raw_value = raw_value[1:] + if raw_value.endswith("\n"): + raw_value = raw_value[:-1] + return raw_value + + +def _coerce_qwen3_coder_parameter( + value: str, + schema: dict[str, Any] | None, +) -> Any: + if not schema: + return value + + typ = schema.get("type") + if isinstance(typ, list): + typ = next((t for t in typ if t != "null"), None) + stripped = value.strip() + + try: + if typ == "boolean": + lowered = stripped.lower() + if lowered == "true": + return True + if lowered == "false": + return False + if typ == "integer": + return int(stripped) + if typ == "number": + return float(stripped) + if typ in {"array", "object"}: + return json.loads(stripped) + except (TypeError, ValueError, json.JSONDecodeError): + return value + + return value + + +def _build_tool_call( + name: str, + arguments: str, + use_responses: bool, +) -> ChatCompletionMessageFunctionToolCall | ResponseFunctionToolCall: + if use_responses: + return ResponseFunctionToolCall( + type="function_call", + id=f"fc-{uuid.uuid4().hex[:24]}", + call_id=f"call_{uuid.uuid4().hex[:24]}", + name=name, + arguments=arguments, + status="completed", + ) + return ChatCompletionMessageFunctionToolCall( + type="function", + id=f"call_{uuid.uuid4().hex[:24]}", + function=Function(name=name, arguments=arguments), + ) + + +def _process_tool_calls_qwen3_coder_xml( + text: str, + tools: list[Any], + finish_reason: str, + use_responses: bool = False, +) -> tuple[ + list[ChatCompletionMessageFunctionToolCall | ResponseFunctionToolCall] | None, + str, + str, +]: + """Parse Qwen3-Coder XML tool calls with ```` tags. + + SGLang's qwen3_coder parser recognizes the tool name in this format but + can return empty arguments for Claude Code style parameters. That makes + Anthropic tool_use.input become ``{}``, so Claude Code rejects calls like + Bash/Read/Write as missing required parameters. + """ + + reasoning_text, content_text = _detect_think_and_return_ori_think( + text, "", "" + ) + arg_schemas = _tool_argument_schemas(tools) + + tool_calls: list[ChatCompletionMessageFunctionToolCall | ResponseFunctionToolCall] + tool_calls = [] + remove_spans: list[tuple[int, int]] = [] + + for tool_match in _QWEN3_CODER_TOOL_CALL_RE.finditer(content_text): + block_calls: list[ + ChatCompletionMessageFunctionToolCall | ResponseFunctionToolCall + ] = [] + body = tool_match.group("body") + for fn_match in _QWEN3_CODER_FUNCTION_RE.finditer(body): + tool_name = fn_match.group("name") + fn_body = fn_match.group("body") + # A parameter value that itself contains a literal "" + # makes the non-greedy regex truncate at the wrong tag, silently + # producing wrong arguments. When opening/closing tags are unbalanced + # the block is ambiguous, so bail out and let the fallback parser run. + if fn_body.count(""): + return None, text, finish_reason + args: dict[str, Any] = {} + for param_match in _QWEN3_CODER_PARAMETER_RE.finditer(fn_body): + param_name = param_match.group("name") + raw_param_value = param_match.group("value") + # A balanced count of tags still hides nested pairs, e.g. a + # value that contains its own "...". + # The non-greedy regex truncates the outer value at the inner + # closing tag, so treat any parameter marker inside the value as + # ambiguous and fall back to the backend parser. + if ( + "" in raw_param_value + ): + return None, text, finish_reason + param_value = _clean_qwen3_coder_parameter(raw_param_value) + args[param_name] = _coerce_qwen3_coder_parameter( + param_value, + arg_schemas.get(tool_name, {}).get(param_name), + ) + + # No parsed arguments means either a genuinely empty block or a + # truncated/malformed one. Skip it so process_tool_calls falls back + # to the sglang parser instead of committing empty arguments. + if not args: + continue + + arguments = json.dumps(args, ensure_ascii=False) + block_calls.append(_build_tool_call(tool_name, arguments, use_responses)) + + if block_calls: + tool_calls.extend(block_calls) + remove_spans.append(tool_match.span()) + + if not tool_calls: + return None, text, finish_reason + + if finish_reason == "stop": + finish_reason = "tool_calls" + + chunks: list[str] = [] + last = 0 + for start, end in remove_spans: + chunks.append(content_text[last:start]) + last = end + chunks.append(content_text[last:]) + cleaned_text = "".join(chunks).replace("<|im_end|>", "") + + return tool_calls, reasoning_text + cleaned_text, finish_reason + + def _process_tool_calls_sglang( text: str, tools: list[Any], @@ -267,6 +463,16 @@ def process_tool_calls( str, ]: """Process tool calls in the response""" + if tool_call_parser == "qwen3_coder": + tool_calls, output_text, finish_reason = _process_tool_calls_qwen3_coder_xml( + text, + tools, + finish_reason, + use_responses, + ) + if tool_calls is not None: + return tool_calls, output_text, finish_reason + try: return _process_tool_calls_sglang( text, diff --git a/areal/infra/launcher/local.py b/areal/infra/launcher/local.py index 2e1dd82b6d..452c089535 100644 --- a/areal/infra/launcher/local.py +++ b/areal/infra/launcher/local.py @@ -31,6 +31,7 @@ JobState, get_scheduling_spec, get_thread_env_vars, + run_post_exit_hook, validate_config_for_launcher, wait_llm_server_addrs, ) @@ -371,7 +372,10 @@ def local_main(config, run_id: int = 0): f"LLM inference server launched at: AREAL_LLM_SERVER_ADDRS={','.join(server_addrs)}" ) except (TimeoutError, KeyboardInterrupt) as e: - launcher.stop_all(signal="SIGINT") + try: + launcher.stop_all(signal="SIGINT") + finally: + run_post_exit_hook(config) raise e if config.get("enable_offload", False): @@ -421,7 +425,10 @@ def local_main(config, run_id: int = 0): remove_status=(), ) except (KeyboardInterrupt, JobException, TimeoutError) as e: - launcher.stop_all("SIGTERM") + try: + launcher.stop_all("SIGTERM") + finally: + run_post_exit_hook(config) # NOTE: For local launcher, We cannot distinguish between a completed job and a failed job. # So we will always try to recover the job if it is finished or failed. recover_states = [JobState.FAILED, JobState.NOT_FOUND, JobState.COMPLETED] diff --git a/areal/infra/launcher/ray.py b/areal/infra/launcher/ray.py index 4b5504221b..6bc7d53bd7 100644 --- a/areal/infra/launcher/ray.py +++ b/areal/infra/launcher/ray.py @@ -34,6 +34,7 @@ JobState, get_scheduling_spec, get_thread_env_vars, + run_post_exit_hook, validate_config_for_distributed_launcher, wait_llm_server_addrs, ) @@ -490,6 +491,7 @@ def sglang_env_hook( launcher.stop_all( force=False ) # force=False will send KeyboardInterrupt to sglang_server.main() to further clean all sglang-related processes + run_post_exit_hook(config) raise e elif allocation_mode.gen_backend == "vllm": config.vllm = to_structured_cfg(config.vllm, vLLMConfig) @@ -531,7 +533,10 @@ def sglang_env_hook( n_vllm_servers, ) except (TimeoutError, KeyboardInterrupt) as e: - launcher.stop_all(force=True) + try: + launcher.stop_all(force=True) + finally: + run_post_exit_hook(config) raise e if config.get("enable_offload", False): @@ -626,6 +631,7 @@ def torch_env_hook(n_tasks: int, placement_group: PlacementGroup) -> list[dict]: # Note: For trainer processes, we use force=True because the trainer doesn't # handle KeyboardInterrupt properly when force=False. launcher.stop_all(force=True, pattern="trainer") + run_post_exit_hook(config) recover_states = [JobState.FAILED] if isinstance(e, JobException): recover_this = ( diff --git a/areal/infra/launcher/slurm.py b/areal/infra/launcher/slurm.py index cb4f1f9428..7f5df3d3b4 100644 --- a/areal/infra/launcher/slurm.py +++ b/areal/infra/launcher/slurm.py @@ -29,6 +29,7 @@ JobState, get_scheduling_spec, get_thread_env_vars, + run_post_exit_hook, validate_config_for_distributed_launcher, wait_llm_server_addrs, ) @@ -571,7 +572,10 @@ def _build_llm_server_plan(backend: str, spec: dict): n_backend_servers, ) except (TimeoutError, KeyboardInterrupt) as e: - launcher.stop_all(force=True) + try: + launcher.stop_all(force=True) + finally: + run_post_exit_hook(config) raise e trainer_n_nodes = n_nodes - n_backend_nodes @@ -675,6 +679,7 @@ def _build_trainer_cmds( ) except (KeyboardInterrupt, JobException, TimeoutError) as e: launcher.stop_all(force=True) + run_post_exit_hook(config) recover_states = [JobState.FAILED, JobState.NOT_FOUND] if isinstance(e, JobException): recover_this = ( diff --git a/areal/infra/rpc/guard/__init__.py b/areal/infra/rpc/guard/__init__.py index af488fbbfd..7c8b721fca 100644 --- a/areal/infra/rpc/guard/__init__.py +++ b/areal/infra/rpc/guard/__init__.py @@ -5,7 +5,7 @@ The Guard is the base process management layer shared between: - ``areal.infra.rpc.rpc_server`` — RPC server (guard + data + engine) -- ``areal.experimental.inference_service.guard`` — inference service guard +- ``areal.v2.inference_service.guard`` — inference service guard Typical usage:: diff --git a/areal/infra/rpc/guard/app.py b/areal/infra/rpc/guard/app.py index 3e13f86dd3..e7067f9de0 100644 --- a/areal/infra/rpc/guard/app.py +++ b/areal/infra/rpc/guard/app.py @@ -5,7 +5,7 @@ This module provides the base Guard functionality shared between: - ``areal.infra.rpc.rpc_server`` (RPC server = guard + data + engine) -- ``areal.experimental.inference_service.guard`` (inference service guard) +- ``areal.v2.inference_service.guard`` (inference service guard) Key components: diff --git a/areal/infra/utils/launcher.py b/areal/infra/utils/launcher.py index 973edff381..4b1e73e3b4 100644 --- a/areal/infra/utils/launcher.py +++ b/areal/infra/utils/launcher.py @@ -5,6 +5,7 @@ import getpass import os import pathlib +import subprocess import sys import time @@ -274,3 +275,58 @@ def validate_config_for_distributed_launcher(config): ) if allocation_mode.gen.tp_size > config.cluster.n_gpus_per_node: raise ValueError("Currently only support vLLM TP size <= #GPUs per node.") + + +def run_post_exit_hook(config) -> None: + """Run the post-exit hook command if configured. + + Executes ``config.post_exit_hook`` as a shell command with a 600-second + timeout. The ``LOG_DIR`` environment variable is injected so the hook + can locate experiment logs. Failures (timeout, non-zero exit, exceptions) + are logged as warnings but never propagate — they must not block recovery. + + Args: + config: Experiment configuration with ``post_exit_hook``, + ``experiment_name``, ``trial_name``, and ``cluster.fileroot``. + """ + hook_cmd = getattr(config, "post_exit_hook", "") + if not hook_cmd or not hook_cmd.strip(): + return + + hook_cmd = hook_cmd.strip() + logger.info(f"Running post-exit hook: {hook_cmd}") + + env = os.environ.copy() + try: + log_dir = os.path.join( + config.cluster.fileroot, + "logs", + getpass.getuser(), + config.experiment_name, + config.trial_name, + ) + env["LOG_DIR"] = log_dir + except Exception: + pass + + try: + result = subprocess.run( + hook_cmd, + shell=True, + timeout=600, + capture_output=True, + text=True, + env=env, + ) + if result.stdout: + logger.info(f"Post-exit hook stdout: {result.stdout.strip()}") + if result.stderr: + logger.warning(f"Post-exit hook stderr: {result.stderr.strip()}") + if result.returncode != 0: + logger.warning(f"Post-exit hook exited with code {result.returncode}") + else: + logger.info("Post-exit hook completed successfully") + except subprocess.TimeoutExpired: + logger.warning(f"Post-exit hook timed out after 600s: {hook_cmd}") + except Exception as e: + logger.warning(f"Post-exit hook failed: {e}") diff --git a/areal/infra/workflow_executor.py b/areal/infra/workflow_executor.py index dcf1eabd12..8d00e68877 100644 --- a/areal/infra/workflow_executor.py +++ b/areal/infra/workflow_executor.py @@ -832,6 +832,80 @@ def _get_dump_dir(self, is_eval: bool) -> str | None: subdir = "eval-rollout" if is_eval else "rollout" return os.path.join(log_path, subdir) + @staticmethod + def _compute_output_versions( + sample_versions: list[int], mask: list[int] + ) -> tuple[int, int, list[list[int]]]: + """Filter versions by loss_mask and compute head, tail, and RLE.""" + output_versions = [v for v, m in zip(sample_versions, mask) if m == 1] + head = min(output_versions) if output_versions else -1 + tail = max(output_versions) if output_versions else -1 + rle: list[list[int]] = [] + for v in output_versions: + if rle and rle[-1][0] == v: + rle[-1][1] += 1 + else: + rle.append([v, 1]) + return head, tail, rle + + @staticmethod + def _split_trajectory_for_dump( + ids: list[int], mask: list[int], tokenizer + ) -> dict[str, Any]: + if len(ids) != len(mask): + raise ValueError(f"ids length {len(ids)} != mask length {len(mask)}") + prompt_end = mask.index(1) if 1 in mask else len(ids) + prompt_text = tokenizer.decode(ids[:prompt_end], skip_special_tokens=False) + completion_text = tokenizer.decode(ids[prompt_end:], skip_special_tokens=False) + + # Count generation runs without decoding to avoid unnecessary work + # in the common single-turn case. + n = len(mask) + gen_count = 0 + idx = 0 + while idx < n: + j = idx + 1 + while j < n and mask[j] == mask[idx]: + j += 1 + if mask[idx] == 1: + gen_count += 1 + idx = j + + # Only decode segments for multi-turn trajectories so single-turn + # keeps the same 2-decode cost as before. + segments: list[dict[str, Any]] | None = None + if gen_count > 1: + raw_segments: list[dict[str, Any]] = [] + idx = 0 + seen_prompt = False + while idx < n: + j = idx + 1 + while j < n and mask[j] == mask[idx]: + j += 1 + if mask[idx] == 1: + role = "gen" + elif not seen_prompt: + role = "prompt" + seen_prompt = True + else: + role = "context" + raw_segments.append( + { + "role": role, + "len": j - idx, + "text": tokenizer.decode(ids[idx:j], skip_special_tokens=False), + } + ) + idx = j + segments = raw_segments + + return { + "prompt_end": prompt_end, + "prompt_text": prompt_text, + "completion_text": completion_text, + "segments": segments, + } + async def _dump_trajectory( self, traj: dict[str, Any] | None, @@ -841,87 +915,100 @@ async def _dump_trajectory( if traj is None: return False, "trajectory is None" - traj = RTensor.localize(traj) - - dump_dir = self._get_dump_dir(is_eval) - if dump_dir is None: - return False, "dump dir is empty" - - tokenizer = self._get_tokenizer() - if tokenizer is None: - return False, "tokenizer not configured" - - # Extract tensors - input_ids = traj.get("input_ids") - rewards = traj.get("rewards") - loss_mask = traj.get("loss_mask") - attention_mask = traj.get("attention_mask") - - if ( - input_ids is None - or rewards is None - or loss_mask is None - or attention_mask is None - ): - return ( - False, - "missing required tensor fields: input_ids, rewards, attention_mask, or loss_mask", - ) + try: + traj = RTensor.localize(traj) - if "versions" not in traj: - self.logger.warning( - "Trajectory missing 'versions' field, defaulting to current inference engine version." - ) - versions = [self.inference_engine.get_version()] - else: - versions = traj["versions"].flatten().tolist() - - tail_version = max(versions) - head_version = min(versions) - # Create versioned directory - version_dir = os.path.join(dump_dir, str(tail_version)) - await aiofiles.os.makedirs(version_dir, exist_ok=True) - - # Handle batched trajectories - batch_size = input_ids.shape[0] - - file_path = os.path.join(version_dir, f"{task_id}.jsonl") - async with aiofiles.open(file_path, "a") as f: - for i in range(batch_size): - seqlen = attention_mask[i].sum().item() - if seqlen == 0: - continue - ids = input_ids[i, :seqlen].tolist() - mask = loss_mask[i, :seqlen].tolist() - # Skip samples with empty completions (all prompt, no completion tokens) - if mask[-1] != 1: - continue + dump_dir = self._get_dump_dir(is_eval) + if dump_dir is None: + return False, "dump dir is empty" - prompt_end = seqlen - sum(mask) - prompt_ids = ids[:prompt_end] - completion_ids = ids[prompt_end:] + tokenizer = self._get_tokenizer() + if tokenizer is None: + return False, "tokenizer not configured" - # Decode to text - prompt_text = tokenizer.decode(prompt_ids, skip_special_tokens=False) - completion_text = tokenizer.decode( - completion_ids, skip_special_tokens=False + # Extract tensors + input_ids = traj.get("input_ids") + rewards = traj.get("rewards") + loss_mask = traj.get("loss_mask") + attention_mask = traj.get("attention_mask") + + if ( + input_ids is None + or rewards is None + or loss_mask is None + or attention_mask is None + ): + return ( + False, + "missing required tensor fields: input_ids, rewards, attention_mask, or loss_mask", + ) + + if "versions" not in traj: + self.logger.warning( + "Trajectory missing 'versions' field, defaulting to current inference engine version." ) + all_versions = None + default_version = self.inference_engine.get_version() + else: + all_versions = traj["versions"] - reward = rewards[i].item() - - record = { - "task_id": task_id, - "sample_idx": i, - "seqlen": seqlen, - "prompt_len": prompt_end, - "head_version": head_version, - "tail_version": tail_version, - "reward": reward, - "prompt": prompt_text, - "completion": completion_text, - } - await f.write(json.dumps(record) + "\n") - return True, "" + global_tail = ( + all_versions.max().item() + if all_versions is not None + else default_version + ) + # Directory is named by batch-global max version (global_tail). + # Individual records may have tail_version <= global_tail. + version_dir = os.path.join(dump_dir, str(global_tail)) + await aiofiles.os.makedirs(version_dir, exist_ok=True) + + # Handle batched trajectories + batch_size = input_ids.shape[0] + + file_path = os.path.join(version_dir, f"{task_id}.jsonl") + async with aiofiles.open(file_path, "a") as f: + for i in range(batch_size): + seqlen = attention_mask[i].sum().item() + if seqlen == 0: + continue + ids = input_ids[i, :seqlen].tolist() + mask = loss_mask[i, :seqlen].tolist() + # Skip samples with empty completions + if mask[-1] != 1: + continue + + if all_versions is not None: + sample_versions = all_versions[i, :seqlen].tolist() + head_version, tail_version, version_rle = ( + self._compute_output_versions(sample_versions, mask) + ) + else: + head_version = tail_version = default_version + version_rle = [[default_version, sum(mask)]] + + split = self._split_trajectory_for_dump(ids, mask, tokenizer) + + reward = rewards[i].item() + + record = { + "task_id": task_id, + "sample_idx": i, + "seqlen": seqlen, + "prompt_len": split["prompt_end"], + "head_version": head_version, + "tail_version": tail_version, + "version_rle": version_rle, + "reward": reward, + "prompt": split["prompt_text"], + "completion": split["completion_text"], + } + if split["segments"] is not None: + record["segments"] = split["segments"] + + await f.write(json.dumps(record) + "\n") + return True, "" + except Exception as e: + return False, f"dump failed: {e}" def initialize(self, logger=None, train_data_parallel_size: int | None = None): """Initialize the workflow executor and start background threads. diff --git a/areal/models/mcore/bailing_moe.py b/areal/models/mcore/bailing_moe.py index 6157edb129..35d4a28b74 100644 --- a/areal/models/mcore/bailing_moe.py +++ b/areal/models/mcore/bailing_moe.py @@ -200,6 +200,14 @@ def hf_to_mcore_config_bailing_moe( # YaRN RoPE scaling parameters (explicitly set to prevent HybridEngine default drift) "mscale": 0.707, "mscale_all_dim": 0.707, + # original_max_position_embeddings: YaRN nests it inside rope_scaling; non-YaRN + # configs may have it at the top level. Fall back to max_position_embeddings. + # mcore 0.16 default (4096) is wrong for long-context models. + "original_max_position_embeddings": ( + rope_scaling.get("original_max_position_embeddings") + or getattr(hf_config, "original_max_position_embeddings", None) + or getattr(hf_config, "max_position_embeddings", 4096) + ), } # MoE-specific parameters @@ -222,6 +230,10 @@ def hf_to_mcore_config_bailing_moe( "moe_router_load_balancing_type": "none", "moe_grouped_gemm": True, "moe_router_dtype": "fp32", + # Auxiliary-loss-free load balancing (DeepSeek V3 style) + router z-loss. + # 0.0 disables bias updates by default; engine config can override. + "moe_router_bias_update_rate": 0.0, + "moe_z_loss_coeff": 3.5e-6, } # Merge all args diff --git a/areal/models/mcore/bailing_moe_bridge.py b/areal/models/mcore/bailing_moe_bridge.py index 3eead5670e..e23e23a093 100644 --- a/areal/models/mcore/bailing_moe_bridge.py +++ b/areal/models/mcore/bailing_moe_bridge.py @@ -174,6 +174,16 @@ def _build_config(self): "factor", 1.0 ), apply_rope_fusion=False, + # original_max_position_embeddings: YaRN nests it in rope_scaling; fall back + # to the top-level field / max_position_embeddings. mcore 0.16 default (4096) + # is wrong for long-context models. + original_max_position_embeddings=( + (getattr(hf_config, "rope_scaling", None) or {}).get( + "original_max_position_embeddings" + ) + or getattr(hf_config, "original_max_position_embeddings", None) + or getattr(hf_config, "max_position_embeddings", 4096) + ), # MoE parameters moe_ffn_hidden_size=getattr(hf_config, "moe_intermediate_size", None), moe_token_dispatcher_type="alltoall", @@ -192,6 +202,9 @@ def _build_config(self): moe_layer_freq=moe_layer_freq, # Other moe_router_dtype="fp32", + # Auxiliary-loss-free load balancing (DeepSeek V3 style) + router z-loss. + moe_router_bias_update_rate=0.0, + moe_z_loss_coeff=3.5e-6, persist_layer_norm=True, bias_activation_fusion=True, bias_dropout_fusion=True, diff --git a/areal/models/mcore/registry.py b/areal/models/mcore/registry.py index 08cd0a11b2..aada02407d 100644 --- a/areal/models/mcore/registry.py +++ b/areal/models/mcore/registry.py @@ -226,6 +226,20 @@ def make_mcore_model( provider.account_for_embedding_in_pipeline_split = False provider.account_for_loss_in_pipeline_split = False + has_mtp = bool(getattr(provider, "mtp_num_layers", None)) + if mcore_config.enable_mtp: + if not has_mtp: + raise ValueError( + "megatron.enable_mtp=True but the model has no MTP layers." + ) + elif has_mtp: + logger.warning( + "Dropping MTP head (mtp_num_layers=%s -> None); not used in RL and not " + "exportable for Qwen3.6. Set megatron.enable_mtp=True to keep it.", + provider.mtp_num_layers, + ) + provider.mtp_num_layers = None + # LoRA params are injected after model materialization and do not carry # Megatron main_grad buffers required by fused grad accumulation kernels. if use_lora: diff --git a/areal/reward/clevr_count_70k.py b/areal/reward/clevr_count_70k.py index 12f63419c6..1d6ba7bfe1 100644 --- a/areal/reward/clevr_count_70k.py +++ b/areal/reward/clevr_count_70k.py @@ -18,15 +18,17 @@ def extract_answer(pred_str, data_name, use_last_number=True): def clevr_count_70k_reward_fn( prompt, completions, prompt_ids, completion_ids, answer, **kwargs ) -> float: - sol = extract_answer(completions, data_name="") # str number - ans = answer - - if sol is None: - return 0.0 - if ans is None: + try: + sol = extract_answer(str(completions), data_name="") # str number + ans = str(answer) + + if not sol or not ans: + return 0.0 + + is_correct = sol.strip() == ans.strip() + if is_correct: + logger.info(f"completions: {completions}, answer: {answer}") + return float(is_correct) + except Exception: + logger.warning("Exception in clevr_count_70k_reward_fn", exc_info=True) return 0.0 - - is_correct = sol.strip() == ans.strip() - if is_correct: - logger.info(f"completions: {completions}, answer: {answer}") - return float(is_correct) diff --git a/areal/trainer/dpo/dpo_engine.py b/areal/trainer/dpo/dpo_engine.py index 89331f18dd..8cd6736b0b 100644 --- a/areal/trainer/dpo/dpo_engine.py +++ b/areal/trainer/dpo/dpo_engine.py @@ -6,15 +6,15 @@ import torch from areal.api import TrainEngine -from areal.experimental.training_service.controller.controller import ( - GatewayTrainController, -) from areal.infra import TrainController from areal.infra.rpc.serialization import serialize_value from areal.utils import logging, stats_tracker from areal.utils.data import batched_call from areal.utils.functional import dpo_pair_logratios, dpo_preference_loss from areal.utils.perf_tracer import trace_perf +from areal.v2.training_service.controller.controller import ( + GatewayTrainController, +) logger = logging.getLogger("DPOEngine") @@ -153,6 +153,8 @@ def compute_dpo_loss( loss_type: str = "sigmoid", vocab_min_logits: torch.Tensor | None = None, vocab_max_logits: torch.Tensor | None = None, + vocab_mean_logits: torch.Tensor | None = None, + vocab_norm_logits: torch.Tensor | None = None, ) -> torch.Tensor: device = logprobs.device cu_seqlens = input_["cu_seqlens"].to(device=device, dtype=torch.long) diff --git a/areal/trainer/ppo/actor.py b/areal/trainer/ppo/actor.py index 07944a31a5..333f23b92f 100644 --- a/areal/trainer/ppo/actor.py +++ b/areal/trainer/ppo/actor.py @@ -7,9 +7,6 @@ from areal.api import TrainEngine from areal.api.cli_args import MicroBatchSpec, PPOActorConfig, RejectionSamplingConfig -from areal.experimental.training_service.controller.controller import ( - GatewayTrainController, -) from areal.infra import TrainController from areal.infra.rpc.serialization import serialize_value from areal.trainer.ppo.stats import infer_token_denominator @@ -27,15 +24,20 @@ from areal.utils.data import ( KLEstimator, Normalization, + TrajBatchMeta, batched_call, split_padded_tensor_dict_into_mb_list, ) from areal.utils.functional import ( + cispo_loss_fn, ppo_actor_loss_fn, reward_overlong_penalty, sapo_loss_fn, ) from areal.utils.perf_tracer import trace_perf +from areal.v2.training_service.controller.controller import ( + GatewayTrainController, +) logger = logging.getLogger("PPOActor") @@ -140,9 +142,11 @@ def _compute_logp(self, data: dict[str, Any]) -> torch.Tensor | None: @trace_perf("ppo_actor.compute_advantages", category="compute") def compute_advantages(self, data: list[dict[str, Any]]) -> list[dict[str, Any]]: - return batched_call(self._compute_advantages, data) + return batched_call(self._compute_advantages, data, pass_meta=True) - def _compute_advantages(self, data: dict[str, Any]) -> dict[str, Any]: + def _compute_advantages( + self, data: dict[str, Any], meta: TrajBatchMeta | None = None + ) -> dict[str, Any]: bs = data["input_ids"].shape[0] max_seqlen = data["input_ids"].shape[1] batch_indices = torch.arange( @@ -169,8 +173,13 @@ def _compute_advantages(self, data: dict[str, Any]) -> dict[str, Any]: reward_score = torch.clip( reward_score, max=self.reward_clip, min=-self.reward_clip ) + # Use actual trajectory group sizes when available so group-level + # normalization handles failed/filtered rollout samples without slicing + # across prompts. Direct calls without batched metadata keep the legacy + # fixed-group-size behavior. + group_sizes = meta.traj_group_sizes if meta is not None else None if self.reward_norm: - reward_score = self.reward_norm(reward_score) + reward_score = self.reward_norm(reward_score, group_sizes=group_sizes) loss_mask = data["loss_mask"].float() loss_mask = torch.roll(loss_mask, shifts=-1, dims=-1) @@ -201,6 +210,16 @@ def _compute_advantages(self, data: dict[str, Any]) -> dict[str, Any]: seq_no_eos_mask = seqlens == attn_mask.shape[1] rewards = -self.kl_ctl * self.kl_estimator(old_logp, ref_logp) kl_rewards = rewards.clone() + step_rewards = data.get("step_rewards") + step_reward_mask = data.get("step_reward_mask") + if step_rewards is not None: + if step_reward_mask is None: + raise ValueError( + "step_reward_mask is required when step_rewards is provided" + ) + rewards = rewards + step_rewards.to(rewards.dtype) * step_reward_mask.to( + rewards.dtype + ) # KL rewards at the next token after eos is zero. rewards[batch_indices, seqlens - 1] = 0 indices = torch.clip(seqlens - 2, min=0) @@ -236,7 +255,9 @@ def _compute_advantages(self, data: dict[str, Any]) -> dict[str, Any]: # Optionally perform advantage normalization. if self.adv_norm is not None: - advantages = self.adv_norm(advantages, loss_mask) + # Use the same actual trajectory group sizes as reward normalization; + # ignored when adv_norm is batch-level. + advantages = self.adv_norm(advantages, loss_mask, group_sizes=group_sizes) # Store data in the dict. data["advantages"] = advantages @@ -295,6 +316,15 @@ def _ppo_update(self, data: dict[str, Any]) -> None: final_reward=data["tot_rewards"], ) stats_tracker.stat(**stats, denominator="n_valid_tokens") + if step_rewards is not None and step_reward_mask is not None: + step_reward_values = step_rewards.float() * step_reward_mask.float() + stats_tracker.stat( + step_reward=step_reward_values, + denominator="n_valid_tokens", + ) + stats_tracker.scalar( + step_reward_events=float(step_reward_mask.sum().item()), + ) prompt_lens = data["attention_mask"].sum(-1) - data["loss_mask"].sum(-1) seq_stats = dict( @@ -358,6 +388,7 @@ def _ppo_update(self, data: dict[str, Any]) -> None: use_sapo_loss=self.config.use_sapo_loss, sapo_tau_pos=self.config.sapo_tau_pos, sapo_tau_neg=self.config.sapo_tau_neg, + use_cispo_loss=self.config.use_cispo_loss, use_decoupled_loss=self.config.use_decoupled_loss, ), loss_weight_fn=lambda x: x["loss_mask"].count_nonzero(), @@ -420,9 +451,12 @@ def grpo_loss_fn( use_sapo_loss: bool = False, sapo_tau_pos: float = 1.0, sapo_tau_neg: float = 1.05, + use_cispo_loss: bool = False, use_decoupled_loss: bool = False, vocab_min_logits: torch.Tensor | None = None, vocab_max_logits: torch.Tensor | None = None, + vocab_mean_logits: torch.Tensor | None = None, + vocab_norm_logits: torch.Tensor | None = None, ): """Loss function for actor step, all inputs should be splitted into pipeline micro batches, returns loss and logging stats.""" @@ -433,6 +467,9 @@ def grpo_loss_fn( entropy = entropy.detach() + if ProxLogpMethod(prox_logp_method) == ProxLogpMethod.REUSE_TRAIN_LOGP: + prox_logp_gt = logprobs.detach() + # Resolve proximal log-probabilities based on method prox_logp = _resolve_proximal_logp( prox_logp_gt=prox_logp_gt, @@ -447,8 +484,30 @@ def grpo_loss_fn( if m2_threshold is not None: loss_mask = _apply_m2po_masking(old_logp, prox_logp, loss_mask, m2_threshold) - # Use SAPO or PPO loss - if use_sapo_loss: + # Use CISPO, SAPO, or PPO loss + if use_cispo_loss: + if use_sapo_loss: + raise ValueError( + "CISPO and SAPO are mutually exclusive surrogates. " + "Set at most one of use_cispo_loss / use_sapo_loss." + ) + if importance_sampling_level != "token": + raise ValueError( + "CISPO only supports importance_sampling_level='token'. " + "Sequence-level (GSPO-style) CISPO has no published surrogate." + ) + loss, stat = cispo_loss_fn( + logprobs=logprobs, + proximal_logprobs=prox_logp, + advantages=advantages, + eps_clip=eps_clip, + eps_clip_higher=eps_clip_higher, + loss_mask=loss_mask, + old_logprobs=old_logp, + rejection_sampling=rejection_sampling, + cu_seqlens=input_data.get("cu_seqlens"), + ) + elif use_sapo_loss: if use_decoupled_loss: raise ValueError( "SAPO is not compatible with `use_decoupled_loss=True`. " diff --git a/areal/trainer/ppo/critic.py b/areal/trainer/ppo/critic.py index a554181bc2..05f83c7441 100644 --- a/areal/trainer/ppo/critic.py +++ b/areal/trainer/ppo/critic.py @@ -7,9 +7,6 @@ from areal.api import TrainEngine from areal.api.cli_args import MicroBatchSpec, PPOCriticConfig -from areal.experimental.training_service.controller.controller import ( - GatewayTrainController, -) from areal.infra import TrainController from areal.infra.rpc.serialization import serialize_value from areal.trainer.ppo.stats import infer_token_denominator @@ -20,6 +17,9 @@ ) from areal.utils.functional import ppo_critic_loss_fn from areal.utils.perf_tracer import trace_perf +from areal.v2.training_service.controller.controller import ( + GatewayTrainController, +) class PPOCritic: diff --git a/areal/trainer/rl_trainer.py b/areal/trainer/rl_trainer.py index 4747530cab..ef89734813 100644 --- a/areal/trainer/rl_trainer.py +++ b/areal/trainer/rl_trainer.py @@ -35,9 +35,6 @@ vLLMConfig, ) from areal.engine import RemoteSGLangEngine, RemotevLLMEngine -from areal.experimental.inference_service.controller.controller import ( - RolloutControllerV2, -) from areal.infra import ( LocalScheduler, RayScheduler, @@ -58,6 +55,9 @@ from areal.utils.recover import RecoverHandler from areal.utils.saver import Saver from areal.utils.stats_logger import StatsLogger +from areal.v2.inference_service.controller.controller import ( + RolloutControllerV2, +) if TYPE_CHECKING: from datasets import Dataset @@ -323,18 +323,30 @@ def __init__( # Proxy worker initialization (lazy, for AgentWorkflow support) self._proxy_started = False - # Prepare weight update meta and connect to inference engine + # Prepare weight update meta and connect to inference engine. + # v2 controllers pick transport from use_lora: LoRA must go through + # disk (P2P transports cannot carry PEFT-wrapped tensors); non-LoRA + # uses awex. v1 keeps the legacy weight_update_mode dispatch. if self.config.actor._version == "v2": - awex_kwargs: dict[str, Any] = {} if config.actor.use_lora: - awex_kwargs.update( - { - "use_lora": config.actor.use_lora, - "lora_name": config.gconfig.lora_name, - "base_model_name": config.actor.path, - } - ) - self.weight_update_meta = WeightUpdateMeta.from_awex(**awex_kwargs) + disk_kwargs: dict[str, Any] = { + "experiment_name": config.experiment_name, + "trial_name": config.trial_name, + "file_root": config.cluster.fileroot, + "name": "default", + "clear_checkpoint_after_load": True, + "use_lora": config.actor.use_lora, + "lora_name": config.gconfig.lora_name, + "base_model_name": config.actor.path, + # Keep enough recent adapter versions for off-policy + # rollouts (max_head_offpolicyness) plus a safety margin; + # older versions are unloaded to bound sglang VRAM and + # avoid the adapter-accumulation hang. + "lora_keep_versions": config.rollout.max_head_offpolicyness + 2, + } + self.weight_update_meta = WeightUpdateMeta.from_disk(**disk_kwargs) + else: + self.weight_update_meta = WeightUpdateMeta.from_awex() elif self.config.actor.weight_update_mode == "disk": disk_kwargs = { "experiment_name": config.experiment_name, @@ -358,7 +370,7 @@ def __init__( ) self.weight_update_meta = WeightUpdateMeta.from_disk(**disk_kwargs) elif self.config.actor.weight_update_mode == "xccl": - # NCCL/XCCL weight update + # NCCL/XCCL weight update (v1 only) xccl_kwargs: dict[str, Any] = { "gen_allocation": self.rollout_alloc, } diff --git a/areal/trainer/rw/rw_engine.py b/areal/trainer/rw/rw_engine.py index d8c22e832c..1e5cbb0acf 100644 --- a/areal/trainer/rw/rw_engine.py +++ b/areal/trainer/rw/rw_engine.py @@ -5,14 +5,14 @@ import torch from areal.api import TrainEngine -from areal.experimental.training_service.controller.controller import ( - GatewayTrainController, -) from areal.infra import TrainController from areal.infra.rpc.serialization import serialize_value from areal.utils import logging, stats_tracker from areal.utils.data import batched_call from areal.utils.perf_tracer import trace_perf +from areal.v2.training_service.controller.controller import ( + GatewayTrainController, +) logger = logging.getLogger("RWEngine") diff --git a/areal/trainer/sft/lm_engine.py b/areal/trainer/sft/lm_engine.py index 6705af0dbd..5e5343bb2d 100644 --- a/areal/trainer/sft/lm_engine.py +++ b/areal/trainer/sft/lm_engine.py @@ -5,14 +5,14 @@ import torch from areal.api import TrainEngine -from areal.experimental.training_service.controller.controller import ( - GatewayTrainController, -) from areal.infra import TrainController from areal.infra.rpc.serialization import serialize_value from areal.utils import stats_tracker from areal.utils.data import batched_call from areal.utils.perf_tracer import trace_perf +from areal.v2.training_service.controller.controller import ( + GatewayTrainController, +) class LMEngine: @@ -84,9 +84,21 @@ def compute_packed_sft_loss( input_: dict[str, Any], vocab_min_logits: torch.Tensor | None = None, vocab_max_logits: torch.Tensor | None = None, + vocab_mean_logits: torch.Tensor | None = None, + vocab_norm_logits: torch.Tensor | None = None, ) -> torch.Tensor: - """Compute SFT loss from logprobs.""" - del entropy # SFT does not use entropy + """Compute SFT loss from logprobs. + + CP NOTE: MegatronEngine reassembles the CP-local per-token scalars + (logprobs / entropy / vocab_*) into full sequences via + ``reassemble_cp_packed_logprobs`` (a 1D all-gather over the CP group) *before* + calling this loss. The vocab reduction along the V dim happens earlier, so the + all-gather only moves per-token scalars, not the full logits — this is the + OOM-avoiding CP path. Consequently the tensors here are already full-sequence + and CP-replicated; the default DP-only reduce at export time yields the correct + global values. We must NOT additionally reduce across CP (that would multiply + sums by cp_size). See ``reassemble_cp_packed_logprobs`` and #1242. + """ cu_seqlens: torch.Tensor = input_["cu_seqlens"] loss_mask = input_["loss_mask"].bool() @@ -103,11 +115,8 @@ def compute_packed_sft_loss( logp = logprobs[cu_seqlens[i] : cu_seqlens[i + 1]] valid_tokens = int(m.count_nonzero().item()) if valid_tokens == 0: - # This is a padded dummy sequence created in `padded_mb_input`. - # When Ulysses SP is enabled, padded inputs are passed into the loss function. - # So we skip it. + # Padded dummy sequence created in `padded_mb_input`; skip it. continue - n_seqs[i] = True seqlogp[i] = torch.where(m, logp.detach(), 0.0).sum() / valid_tokens @@ -120,6 +129,10 @@ def compute_packed_sft_loss( ) stats_tracker.stat(ppl=(-seqlogp).exp().float(), denominator="n_seqs") stats_tracker.stat(loss=-logprobs.detach(), denominator="n_valid_tokens") + stats_tracker.stat( + entropy=torch.where(loss_mask, entropy.detach().float(), 0.0), + denominator="n_valid_tokens", + ) if vocab_min_logits is not None and vocab_max_logits is not None: stats_tracker.stat( @@ -128,4 +141,11 @@ def compute_packed_sft_loss( denominator="n_tokens", ) + if vocab_mean_logits is not None and vocab_norm_logits is not None: + stats_tracker.stat( + vocab_mean_logits=vocab_mean_logits, + vocab_norm_logits=vocab_norm_logits, + denominator="n_tokens", + ) + return loss diff --git a/areal/utils/constants.py b/areal/utils/constants.py index d10a794b22..dc2e5f4838 100644 --- a/areal/utils/constants.py +++ b/areal/utils/constants.py @@ -40,15 +40,19 @@ class ProxLogpMethod(str, Enum): RECOMPUTE: Standard decoupled PPO - recompute via forward pass. LOGLINEAR: Use log-linear approximation (skip forward pass). METRICS: Recompute + compute approximation metrics for evaluation. + REUSE_TRAIN_LOGP: Reuse training forward-pass logprobs as the proximal + logp. Requires ppo_n_minibatches == 1 so the training forward sees + the unchanged policy. """ RECOMPUTE = "recompute" LOGLINEAR = "loglinear" METRICS = "metrics" + REUSE_TRAIN_LOGP = "reuse_train_logp" def skips_forward_pass(self) -> bool: """Return True if this method skips the forward pass (optimization enabled).""" - return self == ProxLogpMethod.LOGLINEAR + return self in (ProxLogpMethod.LOGLINEAR, ProxLogpMethod.REUSE_TRAIN_LOGP) class ProxApproxMethod(str, Enum): @@ -74,6 +78,7 @@ class ProxApproxMethod(str, Enum): PROX_LOGP_METHOD_RECOMPUTE = ProxLogpMethod.RECOMPUTE.value PROX_LOGP_METHOD_LOGLINEAR = ProxLogpMethod.LOGLINEAR.value PROX_LOGP_METHOD_METRICS = ProxLogpMethod.METRICS.value +PROX_LOGP_METHOD_REUSE_TRAIN_LOGP = ProxLogpMethod.REUSE_TRAIN_LOGP.value # List of all valid prox_logp_method values for configuration PROX_LOGP_METHODS_ALL = [m.value for m in ProxLogpMethod] diff --git a/areal/utils/data.py b/areal/utils/data.py index 09368e4ef2..1560742a1a 100644 --- a/areal/utils/data.py +++ b/areal/utils/data.py @@ -401,10 +401,11 @@ def split_batch( def batched_call( - fn: Callable[[dict[str, Any]], Any], + fn: Callable[..., Any], data: list[dict[str, Any]], *, unpack: bool = True, + pass_meta: bool = False, ) -> Any: """Concatenate per-trajectory dicts into one batch, call *fn*, optionally unpack. @@ -421,9 +422,12 @@ def batched_call( unpack : bool If True (default), split the result back into a per-trajectory list via :func:`split_batch`. + pass_meta : bool + If True, call ``fn(batched, meta)`` so functions that need trajectory + metadata can consume it without injecting sentinel keys into the batch. """ batched, meta = concat_batch(data) - result = fn(batched) + result = fn(batched, meta) if pass_meta else fn(batched) if unpack: return split_batch(result, meta) return result @@ -1393,6 +1397,37 @@ def __init__(self, config: NormConfig): self.group_size = config.group_size self.eps = config.eps + def _build_group_slices( + self, bs: int, group_sizes: list[int] | None + ) -> list[slice]: + """Build slices for group-level normalization. + + When ``group_sizes`` (e.g. ``[8, 7, 8, ...]``) is provided it gives + the actual sample count of each trajectory group, handling variable-size + groups that arise when some rollout samples fail / are filtered. A fixed + ``group_size`` slice would otherwise straddle two groups, or leave a tail + of sequences whose std stays 0 → advantage blows up to (reward-mean)/eps. + When *None*, fall back to fixed-``group_size`` slicing. + """ + if group_sizes is not None: + if any(sz <= 0 for sz in group_sizes): + raise ValueError(f"group_sizes must be all positive, got {group_sizes}") + if sum(group_sizes) != bs: + raise ValueError( + f"group_sizes sum ({sum(group_sizes)}) must equal " + f"batch size ({bs}), got {group_sizes}" + ) + slices: list[slice] = [] + offset = 0 + for sz in group_sizes: + slices.append(slice(offset, offset + sz)) + offset += sz + return slices + return [ + slice(i * self.group_size, (i + 1) * self.group_size) + for i in range(bs // self.group_size) + ] + @torch.no_grad() def __call__( self, @@ -1400,6 +1435,7 @@ def __call__( loss_mask: torch.Tensor | None = None, high_precision: bool = True, reduce_group=None, + group_sizes: list[int] | None = None, ) -> torch.Tensor: bs = x.size(0) eps = self.eps @@ -1408,6 +1444,11 @@ def __call__( if loss_mask is not None and loss_mask.sum().item() == 0: return x.float() + # Pre-compute group slices once (variable-size groups via group_sizes). + group_slices = None + if self.mean_level == "group" or self.std_level == "group": + group_slices = self._build_group_slices(bs, group_sizes) + # Step 1: Compute mean if self.mean_level == "batch": mean = self._compute_mean( @@ -1421,16 +1462,17 @@ def __call__( mean = mean.expand_as(x) elif self.mean_level == "group": mean = torch.zeros_like(x) - for i in range(0, bs // self.group_size): - s = slice(i * self.group_size, (i + 1) * self.group_size) + for s in group_slices: xx = x[s] m = loss_mask[s] if loss_mask is not None else None - - # Special case: with group_size=1 and leave_one_out=True, mean should be 0 - if self.group_size == 1 and self.mean_leave1out: - dtype = torch.float64 if high_precision else torch.float32 - group_mean = torch.zeros( - (1, *xx.shape[1:]), dtype=dtype, device=xx.device + group_sz = s.stop - s.start + + # A singleton group has no peer to leave out. Use itself as the + # baseline so leave-one-out normalization outputs zero instead + # of passing the raw reward/advantage through. + if group_sz == 1 and self.mean_leave1out: + group_mean = xx.to( + torch.float64 if high_precision else torch.float32 ) else: group_mean = self._compute_mean( @@ -1465,14 +1507,14 @@ def __call__( std = std.expand_as(x) elif self.std_level == "group": std = torch.zeros_like(x) - for i in range(0, bs // self.group_size): - s = slice(i * self.group_size, (i + 1) * self.group_size) + for s in group_slices: xx = x[s] m = loss_mask[s] if loss_mask is not None else None group_mean_slice = mean[s] # already computed and expanded + group_sz = s.stop - s.start # Special case: with group_size=1 and std_unbiased=True, std should be 1 for numerical stability - if self.group_size == 1 and self.std_unbiased: + if group_sz == 1 and self.std_unbiased: dtype = torch.float64 if high_precision else torch.float32 group_std = torch.ones( (1, *xx.shape[1:]), dtype=dtype, device=xx.device diff --git a/areal/utils/functional/__init__.py b/areal/utils/functional/__init__.py index c91c3ff2b6..b09c6af5f2 100644 --- a/areal/utils/functional/__init__.py +++ b/areal/utils/functional/__init__.py @@ -3,6 +3,7 @@ from areal.utils.functional.functional import ( RejectionSamplingResult, apply_rejection_sampling, + cispo_loss_fn, dpo_pair_logratios, dpo_preference_loss, masked_normalization, @@ -20,6 +21,7 @@ # functional.py "RejectionSamplingResult", "apply_rejection_sampling", + "cispo_loss_fn", "dpo_pair_logratios", "dpo_preference_loss", "masked_normalization", diff --git a/areal/utils/functional/functional.py b/areal/utils/functional/functional.py index 848edd23d3..8aaed7730a 100644 --- a/areal/utils/functional/functional.py +++ b/areal/utils/functional/functional.py @@ -245,6 +245,17 @@ def apply_rejection_sampling( kl_estimator=estimator_name, apply_clamp=False, # Don't clamp; threshold check handles bounds ) + elif config.metric == "binary_kl": + # Bidirectional binary KL divergence (KPop): mask tokens where + # KL(proximal || behave) > upper OR KL(behave || proximal) > upper. + kl_fwd = compute_binary_kl_divergence( + proximal_logprobs.detach(), old_logprobs.detach() + ) + kl_rev = compute_binary_kl_divergence( + old_logprobs.detach(), proximal_logprobs.detach() + ) + # Use max of forward and reverse as the metric; upper bound applies to both. + metric = torch.maximum(kl_fwd, kl_rev) else: raise ValueError(f"Unknown metric: {config.metric}") @@ -426,6 +437,18 @@ def apply_rejection_sampling( ) +def compute_binary_kl_divergence( + log_p: torch.Tensor, log_q: torch.Tensor, eps: float = 1e-8 +) -> torch.Tensor: + """KL(P||Q) for Bernoulli distributions parameterized by log-probabilities. + Treats each element as a Bernoulli: P = [p, 1-p], Q = [q, 1-q]. + KL(P||Q) = p*log(p/q) + (1-p)*log((1-p)/(1-q)) + """ + p = torch.clamp(torch.exp(log_p), eps, 1.0 - eps) + q = torch.clamp(torch.exp(log_q), eps, 1.0 - eps) + return p * torch.log(p / q) + (1 - p) * torch.log((1 - p) / (1 - q)) + + def ppo_actor_loss_fn( logprobs: torch.Tensor, proximal_logprobs: torch.Tensor, @@ -549,6 +572,7 @@ def ppo_actor_loss_fn( clip_mask=clip_mask, dual_clip_mask=dual_clip_mask, ) + if rejection_sampling is not None: stat.update( behave_approx_kl=behave_approx_kl.detach(), @@ -638,6 +662,124 @@ def sapo_loss_fn( return pg_loss, stat +def cispo_loss_fn( + logprobs: torch.Tensor, + proximal_logprobs: torch.Tensor, + advantages: torch.Tensor, + eps_clip: float, + loss_mask: torch.Tensor, + eps_clip_higher: float | None = None, + old_logprobs: torch.Tensor | None = None, + rejection_sampling: RejectionSamplingConfig | None = None, + cu_seqlens: torch.Tensor | None = None, +) -> tuple[torch.Tensor, dict]: + """CISPO (Clipped IS-weight Policy Optimization) loss from MiniMax-M1. + + PPO/GRPO-style clipping zeroes the gradient of any token whose + importance-sampling ratio leaves the clip band: ``min(r*A, clip(r)*A)`` is + constant in theta there. MiniMax-M1 (https://arxiv.org/abs/2506.13585, + Eq. 4-5) observes those are disproportionately low-probability "fork" tokens + (``However``, ``Wait``, ...) that steer reasoning, and instead clips the IS + *weight* under stop-gradient while keeping the gradient on every token's + ``log pi_theta`` (the choice ScaleRL, https://arxiv.org/abs/2510.13786 Eq. 4, + also adopts). Per token:: + + ratio = exp(logprobs - proximal_logprobs) + ratio_clipped = clip(ratio, 1 - eps_clip, 1 + eps_clip_higher) # stop-grad + pg_loss = -sg(ratio_clipped) * advantages * logprobs + + Advantages are never clipped. The clip bounds reuse the same delta-from-1 + convention as :func:`ppo_actor_loss_fn`. CISPO is canonically single-sided: + pass ``eps_clip=1.0`` (lower bound 0) with ``eps_clip_higher=4.0`` for the + wide MiniMax-M1 range. Token level only -- the geometric-mean sequence ratio + of GSPO is not part of the MiniMax-M1 surrogate. + + Decoupled loss: when ``rejection_sampling`` is set, the detached + ``pi_proximal/pi_behave`` weight (from ``old_logprobs``) rescales each token's + surrogate, mirroring :func:`ppo_actor_loss_fn`. + + Args: + logprobs: Current policy log-probabilities (pi_theta), with autograd. + proximal_logprobs: Proximal policy log-probabilities; enters only through + the detached ratio path, so it carries no gradient. + advantages: Per-token advantage estimates; detached on entry. + eps_clip: Lower clipping delta from 1 (ratio lower bound ``1 - eps_clip``). + loss_mask: Mask for valid tokens (1 = valid). + eps_clip_higher: Upper clipping delta from 1 (ratio upper bound + ``1 + eps_clip_higher``); must be positive -- the asymmetric upper + clip is the defining knob of CISPO, so ``None`` is rejected. + old_logprobs: Behavior policy log-probabilities (pi_behave) from inference. + Required only when ``rejection_sampling`` is set (decoupled loss). + rejection_sampling: Staleness filtering / off-policy correction config. + None disables it (pure on-policy CISPO). + cu_seqlens: Cumulative sequence lengths for 1D packed inputs; required when + ``rejection_sampling.level == "sequence"``. + + Returns: + ``(loss, stat)`` matching the PPO loss signature. ``stat['clip_mask']`` + flags tokens whose raw ratio left the band (CISPO never zeroes their + loss, so band-exit -- not loss-affecting clip -- is the meaningful + metric); ``stat['importance_weight']`` reports the unclipped ratio. When + rejection sampling is active the ``behave_*`` keys are reported too. + """ + if eps_clip_higher is None or eps_clip_higher <= 0: + raise ValueError( + "CISPO requires a positive eps_clip_higher; the asymmetric upper " + f"clip is the defining knob (MiniMax-M1 Eq. 4-5). Got {eps_clip_higher!r}." + ) + # Pre-rejection token count, so the denominator matches loss_weight_fn. + loss_mask_count = loss_mask.count_nonzero() or 1 + + # Decoupled off-policy correction: the pi_proximal/pi_behave weight. + if rejection_sampling is not None: + rs_result = apply_rejection_sampling( + proximal_logprobs=proximal_logprobs, + old_logprobs=old_logprobs, + loss_mask=loss_mask, + cu_seqlens=cu_seqlens, + config=rejection_sampling, + ) + loss_mask = rs_result.loss_mask + behave_imp_weight = rs_result.behave_imp_weight + filtered_fraction = rs_result.filtered_fraction + + advantages = advantages.detach() + + # Stop-gradient on the clipped IS weight; gradient flows through logprobs only. + log_ratio = (logprobs - proximal_logprobs).detach() + ratio = torch.exp(log_ratio) + ratio_clipped = torch.clamp(ratio, 1.0 - eps_clip, 1.0 + eps_clip_higher).detach() + pg_loss = -ratio_clipped * advantages * logprobs + + if rejection_sampling is not None: + # behave_imp_weight is detached at source -> still a valid policy gradient. + behave_approx_kl = proximal_logprobs.detach() - old_logprobs.detach() + behave_mask = (behave_imp_weight > 0).logical_and(loss_mask.bool()) + behave_approx_kl = torch.where(behave_mask, behave_approx_kl, 0.0) + pg_loss = pg_loss * behave_imp_weight + + logging_loss = pg_loss.detach() + pg_loss = torch.where(loss_mask, pg_loss, 0).sum() / loss_mask_count + + clip_mask = (ratio_clipped != ratio).logical_and(loss_mask) + stat = dict( + loss=logging_loss, + importance_weight=ratio.detach(), + approx_kl=log_ratio, + clip_mask=clip_mask, + # CISPO has no dual clip; zeros keep the stat schema stable. + dual_clip_mask=torch.zeros_like(loss_mask, dtype=torch.bool), + ) + if rejection_sampling is not None: + stat.update( + behave_approx_kl=behave_approx_kl.detach(), + behave_imp_weight=behave_imp_weight.detach(), + behave_mask=behave_mask, + filtered_fraction=filtered_fraction, + ) + return pg_loss, stat + + def dpo_pair_logratios( logprobs: torch.Tensor, ref_logprobs: torch.Tensor, diff --git a/areal/utils/logging.py b/areal/utils/logging.py index 94e5839a6e..a46e114375 100644 --- a/areal/utils/logging.py +++ b/areal/utils/logging.py @@ -40,10 +40,12 @@ "LocalScheduler": "blue", "RayScheduler": "blue", "SlurmScheduler": "blue", + "InfLocalScheduler": "blue", # Launchers - blue "LocalLauncher": "blue", "RayLauncher": "blue", "SlurmLauncher": "blue", + "InfCli": "blue", # Workflows - purple "RLVRWorkflow": "light_purple", "VisionRLVRWorkflow": "light_purple", @@ -53,6 +55,7 @@ "TrainController": "white", "RolloutController": "white", "WorkflowExecutor": "white", + "AgentCli": "white", # Stats/Perf - green "StatsLogger": "light_green", "StatsTracker": "light_green", @@ -108,6 +111,8 @@ "ToolCallParser": "light_purple", "TokenLogpReward": "light_purple", "ProxyUtils": "light_purple", + "AReaL-SWEAgent": "light_purple", + "SWETrain": "light_green", # Agent Service - purple "AgentGateway": "light_purple", "AgentBridge": "light_purple", diff --git a/areal/utils/network.py b/areal/utils/network.py index 3481016039..8cf9f5955e 100644 --- a/areal/utils/network.py +++ b/areal/utils/network.py @@ -139,8 +139,10 @@ def find_free_ports( free_ports = [] attempted_ports = set() - # Calculate available port range - available_range = max_port - min_port + 1 - len(exclude_ports) + # Calculate available port range. Only excluded ports that fall within + # [min_port, max_port] reduce availability; out-of-range entries do not. + in_range_excluded = sum(min_port <= p <= max_port for p in exclude_ports) + available_range = max_port - min_port + 1 - in_range_excluded if count > available_range: raise ValueError( diff --git a/areal/utils/recover.py b/areal/utils/recover.py index 13d1016cca..9af81b7efb 100644 --- a/areal/utils/recover.py +++ b/areal/utils/recover.py @@ -181,7 +181,7 @@ def _is_gateway_train_controller( | TrainController | dict[str, TrainEngine | TrainController], ) -> bool: - from areal.experimental.training_service.controller.controller import ( + from areal.v2.training_service.controller.controller import ( GatewayTrainController, ) diff --git a/areal/utils/stats_tracker.py b/areal/utils/stats_tracker.py index ec8c9617af..64bb9652b0 100644 --- a/areal/utils/stats_tracker.py +++ b/areal/utils/stats_tracker.py @@ -36,6 +36,11 @@ def __init__(self, name: str = ""): self.scope_stack.append(name.strip("/")) self.denominators = {} # key -> denominator key self.reduce_types = {} # key -> ReduceType + # Per-key override of the reduce_group. If set, its value takes + # precedence over the `reduce_group` passed to `export`. This is used, + # e.g., to make CP-local SFT stats (loss/entropy/vocab_*) reduce across + # DP + CP so the reported numbers are CP-invariant (#1242 follow-up). + self.reduce_groups = {} # key -> dist.ProcessGroup self.stats = defaultdict(list) @@ -93,7 +98,7 @@ def record_timing(self, key): self._set_reduce_type(full_key, ReduceType.SCALAR) self.stats[full_key].append(time.perf_counter() - start_time) - def denominator(self, **kwargs): + def denominator(self, *, reduce_group=None, **kwargs): with self.lock: for key, value in kwargs.items(): if not isinstance(value, torch.Tensor) or value.dtype != torch.bool: @@ -105,21 +110,34 @@ def denominator(self, **kwargs): full_key = self._get_full_key(key) self._set_reduce_type(full_key, ReduceType.SUM) self.stats[full_key].append(value.detach().clone()) + if reduce_group is not None: + self.reduce_groups[full_key] = reduce_group - def scalar(self, **kwargs): + def scalar(self, *, reduce_group=None, **kwargs): with self.lock: for key, value in kwargs.items(): full_key = self._get_full_key(key) self._set_reduce_type(full_key, ReduceType.SCALAR) self.stats[full_key].append(float(value)) + if reduce_group is not None: + self.reduce_groups[full_key] = reduce_group def stat( self, denominator: str, reduce_type: ReduceType | None = None, + *, + reduce_group=None, **kwargs, ): - """Record multiple values from a dictionary""" + """Record multiple values from a dictionary. + + If `reduce_group` is provided, it overrides the `reduce_group` passed + to `export` for these specific keys. This enables recording stats that + must be reduced across a different topology than the default (e.g., + loss/vocab_* under CP-local loss need DP+CP reduce while n_seqs stays + DP-only). + """ with self.lock: for key, value in kwargs.items(): if not isinstance(value, torch.Tensor) or value.dtype != torch.float: @@ -143,12 +161,18 @@ def stat( reduce_type = ReduceType.AVG_MIN_MAX self._set_reduce_type(full_key, reduce_type) self.stats[full_key].append(value.detach().clone()) + if reduce_group is not None: + self.reduce_groups[full_key] = reduce_group def _set_reduce_type(self, key, reduce_type): if not isinstance(reduce_type, ReduceType): raise ValueError("reduce_type must be a ReduceType enum") self.reduce_types[key] = reduce_type + def _effective_reduce_group(self, key, default_reduce_group): + """Return the per-key reduce_group override if any, else the default.""" + return self.reduce_groups.get(key, default_reduce_group) + def export(self, key=None, reduce_group=None, reset=True) -> dict[str, float]: """Get aggregated statistics""" with self.lock: @@ -159,7 +183,9 @@ def export(self, key=None, reduce_group=None, reset=True) -> dict[str, float]: if full_key in self.denominators: self.denominators.pop(full_key) if full_key in self.reduce_types: - self.denominators.pop(full_key) + self.reduce_types.pop(full_key) + if full_key in self.reduce_groups: + self.reduce_groups.pop(full_key) self.stats.pop(full_key) return result @@ -176,6 +202,7 @@ def export(self, key=None, reduce_group=None, reset=True) -> dict[str, float]: if reset: self.denominators = {} self.reduce_types = {} + self.reduce_groups = {} self.stats = defaultdict(list) results = { k: v.cpu().item() if torch.is_tensor(v) else v @@ -208,9 +235,10 @@ def _aggregate(self, key, reduce_group): sum(self.stats[key]), dtype=torch.float32, device=device ) cnt = torch.tensor(len(self.stats[key]), dtype=torch.float32, device=device) - if reduce_group is not None: - dist.all_reduce(value, group=reduce_group) - dist.all_reduce(cnt, group=reduce_group) + effective_group = self._effective_reduce_group(key, reduce_group) + if effective_group is not None: + dist.all_reduce(value, group=effective_group) + dist.all_reduce(cnt, group=effective_group) result[key] = float(value / cnt) result[key + "__count"] = int(cnt) else: @@ -223,10 +251,11 @@ def _aggregate(self, key, reduce_group): def _sum_of(self, key, reduce_group): values = self.stats[key] + effective_group = self._effective_reduce_group(key, reduce_group) if key not in self.denominators: x = sum([x.sum() for x in values]) - if reduce_group is not None: - dist.all_reduce(x, group=reduce_group) + if effective_group is not None: + dist.all_reduce(x, group=effective_group) else: denominator = self.denominators[key] if denominator not in self.stats: @@ -237,8 +266,8 @@ def _sum_of(self, key, reduce_group): for v, d in zip(values, self.stats[denominator]): xs.append(torch.where(d, v, 0.0).sum()) x = sum(xs) - if reduce_group is not None: - dist.all_reduce(x, group=reduce_group) + if effective_group is not None: + dist.all_reduce(x, group=effective_group) return float(x) def _avg_of(self, key, reduce_group): @@ -253,9 +282,10 @@ def _avg_of(self, key, reduce_group): ds.append(d.sum()) x = sum(xs) d = sum(ds) - if reduce_group is not None: - dist.all_reduce(x, group=reduce_group) - dist.all_reduce(d, group=reduce_group) + effective_group = self._effective_reduce_group(key, reduce_group) + if effective_group is not None: + dist.all_reduce(x, group=effective_group) + dist.all_reduce(d, group=effective_group) if d == 0: return None return x / d @@ -269,8 +299,9 @@ def _min_of(self, key, reduce_group): for v, d in zip(values, self.stats[denominator]): xs.append(torch.where(d, v, float("inf")).min()) x = min(xs) - if reduce_group is not None: - dist.all_reduce(x, group=reduce_group, op=dist.ReduceOp.MIN) + effective_group = self._effective_reduce_group(key, reduce_group) + if effective_group is not None: + dist.all_reduce(x, group=effective_group, op=dist.ReduceOp.MIN) if torch.isinf(x): return None return float(x) @@ -284,8 +315,9 @@ def _max_of(self, key, reduce_group): for v, d in zip(values, self.stats[denominator]): xs.append(torch.where(d, v, -float("inf")).max()) x = max(xs) - if reduce_group is not None: - dist.all_reduce(x, group=reduce_group, op=dist.ReduceOp.MAX) + effective_group = self._effective_reduce_group(key, reduce_group) + if effective_group is not None: + dist.all_reduce(x, group=effective_group, op=dist.ReduceOp.MAX) if torch.isinf(x): return None return float(x) diff --git a/areal/experimental/inference_service/data_proxy/__init__.py b/areal/v2/__init__.py similarity index 100% rename from areal/experimental/inference_service/data_proxy/__init__.py rename to areal/v2/__init__.py diff --git a/areal/experimental/agent_service/README.md b/areal/v2/agent_service/README.md similarity index 87% rename from areal/experimental/agent_service/README.md rename to areal/v2/agent_service/README.md index 30f76bb8ec..de8f32d80a 100644 --- a/areal/experimental/agent_service/README.md +++ b/areal/v2/agent_service/README.md @@ -32,7 +32,8 @@ Client (HTTP/WS) ### Components **Gateway** — Public entry point. Accepts WebSocket connections (Gateway protocol) and -HTTP requests (OpenResponses bridge at `POST /v1/responses`). Routes to the appropriate +HTTP requests via two bridges: the OpenResponses bridge (`POST /v1/responses`) and the +OpenAI chat-completions bridge (`POST /v1/chat/completions`). Routes to the appropriate DataProxy via the Router. **Router** — Session-affine routing service. DataProxy instances register at startup. @@ -123,11 +124,12 @@ class EventEmitter(Protocol): ### Gateway -| Endpoint | Method | Description | -| --------------- | ------ | -------------------------- | -| `/health` | GET | Health check | -| `/ws` | WS | Gateway WebSocket protocol | -| `/v1/responses` | POST | OpenResponses HTTP bridge | +| Endpoint | Method | Description | +| ---------------------- | ------ | ------------------------------ | +| `/health` | GET | Health check | +| `/ws` | WS | Gateway WebSocket protocol | +| `/v1/responses` | POST | OpenResponses HTTP bridge | +| `/v1/chat/completions` | POST | OpenAI chat-completions bridge | ## Multi-turn Conversation Flow @@ -152,7 +154,7 @@ Turn 2: ## Code Organization ``` -areal/experimental/agent_service/ +areal/v2/agent_service/ ├── __init__.py # Public exports (AgentRequest, AgentResponse, etc.) ├── README.md # This document ├── auth.py # Admin key auth helpers (hmac-safe comparison) @@ -163,29 +165,29 @@ areal/experimental/agent_service/ │ └── controller.py # AgentController orchestrator ├── guard/ │ ├── __init__.py # Module docstring -│ ├── __main__.py # python -m areal.experimental.agent_service.guard +│ ├── __main__.py # python -m areal.v2.agent_service.guard │ └── app.py # Guard Flask app (pass-through to areal.infra.rpc.guard) ├── gateway/ │ ├── __init__.py # Public exports -│ ├── __main__.py # python -m areal.experimental.agent_service.gateway +│ ├── __main__.py # python -m areal.v2.agent_service.gateway │ ├── app.py # create_gateway_app() │ ├── bridge.py # OpenResponsesBridge, mount_bridge() │ └── config.py # GatewayConfig dataclass ├── router/ │ ├── __init__.py # Public exports -│ ├── __main__.py # python -m areal.experimental.agent_service.router +│ ├── __main__.py # python -m areal.v2.agent_service.router │ ├── app.py # create_router_app() │ ├── client.py # RouterClient │ └── config.py # RouterConfig dataclass ├── data_proxy/ │ ├── __init__.py # Public exports -│ ├── __main__.py # python -m areal.experimental.agent_service.data_proxy +│ ├── __main__.py # python -m areal.v2.agent_service.data_proxy │ ├── app.py # create_data_proxy_app() │ ├── client.py # DataProxyClient │ └── config.py # DataProxyConfig dataclass └── worker/ ├── __init__.py # Public exports - ├── __main__.py # python -m areal.experimental.agent_service.worker + ├── __main__.py # python -m areal.v2.agent_service.worker ├── app.py # create_worker_app() └── config.py # WorkerConfig dataclass diff --git a/areal/experimental/agent_service/__init__.py b/areal/v2/agent_service/__init__.py similarity index 87% rename from areal/experimental/agent_service/__init__.py rename to areal/v2/agent_service/__init__.py index 2c964e550f..c92956aeb9 100644 --- a/areal/experimental/agent_service/__init__.py +++ b/areal/v2/agent_service/__init__.py @@ -16,7 +16,12 @@ - ``protocol`` — WebSocket frame types and helpers """ -from .types import AgentRequest, AgentResponse, AgentRunnable, EventEmitter +from .types import ( + AgentRequest, + AgentResponse, + AgentRunnable, + EventEmitter, +) __all__ = [ "AgentRequest", diff --git a/areal/experimental/agent_service/auth.py b/areal/v2/agent_service/auth.py similarity index 100% rename from areal/experimental/agent_service/auth.py rename to areal/v2/agent_service/auth.py diff --git a/areal/experimental/agent_service/controller/__init__.py b/areal/v2/agent_service/controller/__init__.py similarity index 100% rename from areal/experimental/agent_service/controller/__init__.py rename to areal/v2/agent_service/controller/__init__.py diff --git a/areal/experimental/agent_service/controller/controller.py b/areal/v2/agent_service/controller/controller.py similarity index 97% rename from areal/experimental/agent_service/controller/controller.py rename to areal/v2/agent_service/controller/controller.py index f63d08d780..409fd5870c 100644 --- a/areal/experimental/agent_service/controller/controller.py +++ b/areal/v2/agent_service/controller/controller.py @@ -3,7 +3,7 @@ """AgentController — orchestrates agent service micro-services via Guards. Mirrors the architecture of -:class:`~areal.experimental.inference_service.controller.controller.RolloutControllerV2`: +:class:`~areal.v2.inference_service.controller.controller.RolloutControllerV2`: Guard workers are created via the Scheduler, then the controller forks Router, Worker+DataProxy pairs, and Gateway onto them via HTTP API. @@ -126,7 +126,7 @@ def _do_initialize(self) -> None: # Step 1: Create Guard workers via scheduler guard_spec = SchedulingSpec(**asdict(cfg.scheduling_spec[0])) guard_spec.gpu = 0 - guard_spec.cmd = f"{sys.executable} -m areal.experimental.agent_service.guard" + guard_spec.cmd = f"{sys.executable} -m areal.v2.agent_service.guard" self._base_env = dict(guard_spec.env_vars) num_guards = _DEFAULT_PAIR_COUNT guard_job = Job( @@ -150,7 +150,7 @@ def _do_initialize(self) -> None: router_cmd = [ sys.executable, "-m", - "areal.experimental.agent_service.router", + "areal.v2.agent_service.router", "--admin-api-key", cfg.admin_api_key, ] @@ -170,7 +170,7 @@ def _do_initialize(self) -> None: gw_cmd = [ sys.executable, "-m", - "areal.experimental.agent_service.gateway", + "areal.v2.agent_service.gateway", "--router-addr", self._router_addr, "--admin-api-key", @@ -246,7 +246,7 @@ def scale_up(self, count: int) -> list[int]: worker_cmd = [ sys.executable, "-m", - "areal.experimental.agent_service.worker", + "areal.v2.agent_service.worker", "--agent", cfg.agent_cls_path, "--log-level", @@ -263,7 +263,7 @@ def scale_up(self, count: int) -> list[int]: proxy_cmd = [ sys.executable, "-m", - "areal.experimental.agent_service.data_proxy", + "areal.v2.agent_service.data_proxy", "--worker-addr", worker_addr, ] diff --git a/areal/experimental/agent_service/data_proxy/__init__.py b/areal/v2/agent_service/data_proxy/__init__.py similarity index 100% rename from areal/experimental/agent_service/data_proxy/__init__.py rename to areal/v2/agent_service/data_proxy/__init__.py diff --git a/areal/experimental/agent_service/data_proxy/__main__.py b/areal/v2/agent_service/data_proxy/__main__.py similarity index 94% rename from areal/experimental/agent_service/data_proxy/__main__.py rename to areal/v2/agent_service/data_proxy/__main__.py index cff27a6816..26c5b3a9a1 100644 --- a/areal/experimental/agent_service/data_proxy/__main__.py +++ b/areal/v2/agent_service/data_proxy/__main__.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 -"""``python -m areal.experimental.agent_service.data_proxy``""" +"""``python -m areal.v2.agent_service.data_proxy``""" import argparse diff --git a/areal/v2/agent_service/data_proxy/app.py b/areal/v2/agent_service/data_proxy/app.py new file mode 100644 index 0000000000..b5120f756c --- /dev/null +++ b/areal/v2/agent_service/data_proxy/app.py @@ -0,0 +1,283 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Data Proxy — stateful session proxy between Gateway and Worker.""" + +from __future__ import annotations + +import asyncio +import time +from dataclasses import dataclass, field +from typing import Any + +import httpx +from fastapi import FastAPI, HTTPException +from fastapi.responses import JSONResponse, StreamingResponse + +from areal.utils import logging + +from ..protocol import PASSTHROUGH_HEADER +from .config import DataProxyConfig + +logger = logging.getLogger("AgentDataProxy") + + +@dataclass +class _SessionData: + history: list[dict[str, Any]] = field(default_factory=list) + metadata: dict[str, Any] = field(default_factory=dict) + last_active: float = field(default_factory=time.monotonic) + reward: float | None = None + # Per-session inference routing for self-evolution. Holds + # ``{"base_url", "api_key", "model"}`` where ``api_key`` is the + # ``sk-sess-*`` the **caller** obtained itself and passed on the turn + # (``session_api_key``). The Agent Service never talks to the training + # side — it only forwards these fields to the worker so the agent routes + # its LLM calls through the inference gateway under that key. Cached on + # the first turn that carries them so later turns of a multi-turn session + # can omit them. + inference: dict[str, Any] | None = None + + +def create_data_proxy_app(config: DataProxyConfig) -> FastAPI: + app = FastAPI(title="AReaL Data Proxy") + sessions: dict[str, _SessionData] = {} + http_client = httpx.AsyncClient(timeout=config.request_timeout) + + async def _close_worker_session(session_key: str) -> None: + try: + await http_client.post( + f"{config.worker_addr}/session/{session_key}/close", timeout=5 + ) + except Exception: + logger.debug("Failed to close worker session %s", session_key) + + def _resolve_inference( + session: _SessionData, body: dict[str, Any], metadata: dict[str, Any] + ) -> dict[str, Any]: + """Forward caller-supplied inference routing to the worker metadata. + + Self-evolution decouples the Agent Service from the training side: the + **caller** mints its own per-session ``sk-sess-*`` (e.g. via its own + ``/rl/start_session``) and passes it on the turn body. This proxy never + contacts the inference/training side — it merely caches the routing + handle on the session and injects it as ``metadata['areal_inference']`` + so the agent routes its LLM calls through the inference gateway. + + The turn opts in **by the presence of the routing fields** (no separate + flag); the required pair is: + + ``inf_base_url`` — inference gateway base URL the agent's LLM + calls go to (required). + ``session_api_key`` — the caller-minted ``sk-sess-*`` (required). + ``inf_model`` — model id the agent should request (default ""). + + ``inf_model`` is optional and never triggers self-evolution on its own; + only the required pair does. The handle is cached on the first turn that + carries it, so a multi-turn session may send these fields once and omit + them afterwards. + """ + base_url = (body.get("inf_base_url") or "").rstrip("/") + api_key = body.get("session_api_key") or "" + if base_url or api_key: + if not base_url or not api_key: + raise HTTPException( + status_code=400, + detail=( + "self-evolution requires both 'inf_base_url' and " + "'session_api_key' in the turn body" + ), + ) + session.inference = { + "base_url": base_url, + "api_key": api_key, + "model": body.get("inf_model", "") or "", + } + + if session.inference is not None: + metadata = {**metadata, "areal_inference": dict(session.inference)} + return metadata + + async def _reap_idle_sessions() -> None: + while True: + await asyncio.sleep(60) + now = time.monotonic() + stale = [ + k + for k, s in sessions.items() + if now - s.last_active > config.session_timeout + ] + for k in stale: + del sessions[k] + await _close_worker_session(k) + if stale: + logger.info("Reaped %d idle sessions", len(stale)) + + @app.on_event("startup") + async def startup(): + app.state.reaper_task = asyncio.create_task(_reap_idle_sessions()) + + @app.on_event("shutdown") + async def shutdown(): + await http_client.aclose() + + @app.get("/health") + async def health(): + return { + "status": "ok", + "active_sessions": len(sessions), + "worker_addr": config.worker_addr, + } + + @app.post("/session/{session_key}/turn") + async def turn(session_key: str, body: dict[str, Any]): + """Single turn endpoint for every protocol and streaming mode. + + The worker's ``/run`` decides the shape of the turn and signals it via + the :data:`PASSTHROUGH_HEADER` response header: + + - header absent — a structured turn (``application/json``). The body is + read fully, conversation history is rebuilt from the emitted + ``events``, and the JSON is returned to the caller (this backs + ``/v1/responses`` and the WebSocket path). + - header == ``"1"`` — a raw-passthrough turn. The body is relayed + **byte-for-byte** without parsing, so the caller gets the upstream's + exact wire format (this backs ``/v1/chat/completions``, streaming or + not). Keying on the marker rather than ``Content-Type`` means a + *non-streaming* passthrough — itself ``application/json`` — is still + relayed verbatim instead of being mistaken for a structured turn. No + history is kept on this path; stateful callers rely on *route + affinity* (a stable ``session_key`` pins every turn to this same + DataProxy/Worker so the agent reuses its own state). + """ + session = sessions.get(session_key) + if session is None: + session = _SessionData() + sessions[session_key] = session + + message = body.get("message", "") + run_id = body.get("run_id", "") + queue_mode = body.get("queue_mode", "collect") + metadata = body.get("metadata", {}) + + # Self-evolution (opt-in): when the turn carries inference-routing + # fields, forward them so the agent's LLM calls flow through the + # inference gateway under the caller's own ``sk-sess-*``. No-op for a + # plain turn that omits them. + metadata = _resolve_inference(session, body, metadata) + + worker_request = { + "message": message, + "session_key": session_key, + "run_id": run_id, + "history": session.history.copy(), + "queue_mode": queue_mode, + "metadata": metadata, + } + + # Open the worker stream so we can inspect status/Content-Type before + # deciding whether to parse (structured) or relay (raw passthrough). + req = http_client.build_request( + "POST", f"{config.worker_addr}/run", json=worker_request + ) + resp = await http_client.send(req, stream=True) + session.last_active = time.monotonic() + + is_passthrough = resp.headers.get(PASSTHROUGH_HEADER) == "1" + + if is_passthrough: + # Raw-passthrough turn: relay the worker body verbatim, keep no + # history (parsing the stream would fight the byte-exact relay). + async def _relay(): + try: + async for chunk in resp.aiter_raw(): + yield chunk + finally: + await resp.aclose() + session.last_active = time.monotonic() + + headers = { + k: v + for k, v in resp.headers.items() + if k.lower() + not in ( + "content-length", + "transfer-encoding", + "connection", + PASSTHROUGH_HEADER, + ) + } + return StreamingResponse( + _relay(), + status_code=resp.status_code, + headers=headers, + media_type=resp.headers.get("content-type") or None, + ) + + # Structured turn: read the full JSON body, then rebuild history. + await resp.aread() + status_code = resp.status_code + result = resp.json() + await resp.aclose() + + if status_code >= 400: + # Worker reported an error; forward it without touching history. + return JSONResponse(result, status_code=status_code) + + session.history.append({"role": "user", "content": message}) + + call_counter = 0 + for evt in result.get("events", []): + if evt.get("type") == "tool_call": + call_id = f"call_{evt.get('name', '')}_{run_id}_{call_counter}" + call_counter += 1 + session.history.append( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": call_id, + "type": "function", + "function": { + "name": evt.get("name", ""), + "arguments": evt.get("args", ""), + }, + } + ], + } + ) + elif evt.get("type") == "tool_result": + result_call_id = ( + f"call_{evt.get('name', '')}_{run_id}_{call_counter - 1}" + if call_counter > 0 + else f"call_{evt.get('name', '')}_{run_id}_0" + ) + session.history.append( + { + "role": "tool", + "tool_call_id": result_call_id, + "content": evt.get("result", ""), + } + ) + + summary = result.get("summary", "") + if summary: + session.history.append({"role": "assistant", "content": summary}) + + session.last_active = time.monotonic() + return JSONResponse(result, status_code=status_code) + + @app.post("/session/{session_key}/close") + async def close_session(session_key: str): + sessions.pop(session_key, None) + await _close_worker_session(session_key) + return {"status": "ok"} + + @app.get("/session/{session_key}/history") + async def get_history(session_key: str): + session = sessions.get(session_key) + if session is None: + return {"history": []} + return {"history": session.history} + + return app diff --git a/areal/experimental/agent_service/data_proxy/client.py b/areal/v2/agent_service/data_proxy/client.py similarity index 100% rename from areal/experimental/agent_service/data_proxy/client.py rename to areal/v2/agent_service/data_proxy/client.py diff --git a/areal/experimental/agent_service/data_proxy/config.py b/areal/v2/agent_service/data_proxy/config.py similarity index 100% rename from areal/experimental/agent_service/data_proxy/config.py rename to areal/v2/agent_service/data_proxy/config.py diff --git a/areal/experimental/agent_service/gateway/__init__.py b/areal/v2/agent_service/gateway/__init__.py similarity index 100% rename from areal/experimental/agent_service/gateway/__init__.py rename to areal/v2/agent_service/gateway/__init__.py diff --git a/areal/experimental/agent_service/gateway/__main__.py b/areal/v2/agent_service/gateway/__main__.py similarity index 83% rename from areal/experimental/agent_service/gateway/__main__.py rename to areal/v2/agent_service/gateway/__main__.py index 72313bf641..1261112aeb 100644 --- a/areal/experimental/agent_service/gateway/__main__.py +++ b/areal/v2/agent_service/gateway/__main__.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 -"""``python -m areal.experimental.agent_service.gateway``""" +"""``python -m areal.v2.agent_service.gateway``""" import argparse @@ -10,7 +10,12 @@ from ..auth import DEFAULT_ADMIN_API_KEY from .app import create_gateway_app -from .bridge import OpenResponsesBridge, mount_bridge +from .bridge import ( + ChatCompletionsBridge, + OpenResponsesBridge, + mount_bridge, + mount_chat_bridge, +) from .config import GatewayConfig @@ -48,6 +53,12 @@ def main() -> None: ), admin_api_key=config.admin_api_key, ) + mount_chat_bridge( + app, + ChatCompletionsBridge( + router_addr=config.router_addr, admin_api_key=config.admin_api_key + ), + ) uvicorn.run( app, host=config.host, diff --git a/areal/experimental/agent_service/gateway/app.py b/areal/v2/agent_service/gateway/app.py similarity index 81% rename from areal/experimental/agent_service/gateway/app.py rename to areal/v2/agent_service/gateway/app.py index b3043ad57a..88914fedc3 100644 --- a/areal/experimental/agent_service/gateway/app.py +++ b/areal/v2/agent_service/gateway/app.py @@ -7,13 +7,21 @@ import hmac import json import traceback +from typing import Any import httpx -from fastapi import FastAPI, Query, WebSocket, WebSocketDisconnect +from fastapi import ( + Depends, + FastAPI, + HTTPException, + Query, + WebSocket, + WebSocketDisconnect, +) from areal.utils import logging -from ..auth import admin_headers +from ..auth import admin_headers, make_admin_dependency from ..protocol import ( FrameType, RequestFrame, @@ -47,6 +55,7 @@ def create_gateway_app(config: GatewayConfig) -> FastAPI: app = FastAPI(title="AReaL Agent Gateway") http_client = httpx.AsyncClient(timeout=config.forward_timeout) _auth_headers = admin_headers(config.admin_api_key) + _admin = make_admin_dependency(config.admin_api_key) async def _route(session_key: str) -> str: resp = await http_client.post( @@ -82,6 +91,26 @@ async def _send_turn( async def health(): return {"status": "ok"} + @app.post("/sessions/close", dependencies=[Depends(_admin)]) + async def close_session(body: dict[str, Any]): + """Close a session, forwarding to the DataProxy that owns it. + + The ``session_key`` is supplied in the request body. Session lifecycle + is keyed solely by ``session_key`` and is identical across every gateway + surface (WebSocket, ``/v1/responses``, ``/v1/chat/completions``): the + Router pins a key to one DataProxy via route affinity, so resolving the + key here lands on the DataProxy holding that session's state. The + DataProxy close is idempotent and propagates down to the Worker, so an + unknown key is a no-op success. + """ + session_key = body.get("session_key", "") + if not session_key: + raise HTTPException(status_code=400, detail="'session_key' is required") + data_proxy_addr = await _route(session_key) + resp = await http_client.post(f"{data_proxy_addr}/session/{session_key}/close") + resp.raise_for_status() + return {"status": "ok"} + @app.websocket("/ws") async def websocket_endpoint(websocket: WebSocket, token: str = Query(default="")): if not hmac.compare_digest(token, config.admin_api_key): diff --git a/areal/v2/agent_service/gateway/bridge.py b/areal/v2/agent_service/gateway/bridge.py new file mode 100644 index 0000000000..0d5936b5a2 --- /dev/null +++ b/areal/v2/agent_service/gateway/bridge.py @@ -0,0 +1,487 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""OpenResponses HTTP bridge — translates POST /v1/responses to DataProxy turns.""" + +from __future__ import annotations + +import json +import uuid +from abc import ABC, abstractmethod +from collections.abc import AsyncIterator +from typing import Any + +import httpx +from fastapi import Depends, FastAPI, Request +from fastapi.responses import JSONResponse, StreamingResponse + +from areal.utils import logging + +from ..auth import DEFAULT_ADMIN_API_KEY, admin_headers, make_admin_dependency +from ..protocol import generate_run_id + +logger = logging.getLogger("AgentBridge") + +# Request/response header carrying the session identifier. A client sends it to +# pin a multi-turn conversation onto the same DataProxy/Worker (route affinity); +# both bridges always echo the resolved key back on the response so the client +# can reuse it on the next request — including a key derived or randomly minted +# server-side when the client sent none. +SESSION_KEY_HEADER = "X-AReaL-Session-Key" + + +class AgentBridge(ABC): + @abstractmethod + async def handle_request(self, request: Request) -> Any: ... + + +class OpenResponsesBridge(AgentBridge): + def __init__( + self, router_addr: str, admin_api_key: str = DEFAULT_ADMIN_API_KEY + ) -> None: + self._router_addr = router_addr + self._auth_headers = admin_headers(admin_api_key) + self._http = httpx.AsyncClient(timeout=600.0) + + async def close(self) -> None: + await self._http.aclose() + + async def handle_request(self, request: Request) -> Any: + body = await request.json() + + input_items: list[dict[str, Any]] = body.get("input", []) + instructions: str = body.get("instructions", "") + model: str = body.get("model", "") + user: str = body.get("user", "") + + # An explicit X-AReaL-Session-Key pins the conversation to one + # DataProxy/Worker directly; only when it is absent do we fall back to + # deriving the key from ``user`` (which is then required for affinity). + explicit_session_key = request.headers.get(SESSION_KEY_HEADER) + if not explicit_session_key and not user: + return JSONResponse( + { + "error": { + "message": ( + "either the 'X-AReaL-Session-Key' header or the " + "'user' field is required for session affinity" + ), + "type": "invalid_request", + } + }, + status_code=400, + ) + + message = self._extract_message(input_items, instructions) + session_key = explicit_session_key or self._derive_session_key(user, model) + run_id = generate_run_id() + response_id = f"resp-{uuid.uuid4().hex[:12]}" + + metadata = { + "input": input_items, + "instructions": instructions, + "tools": body.get("tools", []), + "model": model, + "idempotencyKey": response_id, + **body.get("metadata", {}), + } + + try: + route_resp = await self._http.post( + f"{self._router_addr}/route", + json={"session_key": session_key}, + headers=self._auth_headers, + ) + route_resp.raise_for_status() + data_proxy_addr = route_resp.json()["data_proxy_addr"] + + turn_body: dict[str, Any] = { + "message": message, + "run_id": run_id, + "queue_mode": "collect", + "metadata": metadata, + } + # Opt-in self-evolution: forward the caller-supplied inference + # routing fields when present so the DataProxy hands the agent a + # ``sk-sess-*`` (which the caller minted itself) and the agent's LLM + # calls flow through AReaL's inference service. The DataProxy opts + # the turn in by the presence of these fields; absent them the turn + # is plain. The Agent Service never contacts the training side. + for key in ( + "inf_base_url", + "inf_model", + "session_api_key", + ): + if key in body: + turn_body[key] = body[key] + + turn_resp = await self._http.post( + f"{data_proxy_addr}/session/{session_key}/turn", + json=turn_body, + ) + turn_resp.raise_for_status() + result = turn_resp.json() + + output_items = self._build_output_items(result) + response_metadata = result.get("metadata", {}) + + if body.get("stream"): + # Streaming /v1/responses: the agent ran the structured turn to + # completion (collect); re-encode the collected output as an + # OpenAI Responses-format SSE event stream. + return StreamingResponse( + self._responses_sse( + response_id, model, output_items, response_metadata + ), + media_type="text/event-stream", + headers={SESSION_KEY_HEADER: session_key}, + ) + + return JSONResponse( + { + "id": response_id, + "object": "response", + "status": "completed", + "output": output_items, + "model": model, + "metadata": response_metadata, + }, + headers={SESSION_KEY_HEADER: session_key}, + ) + except Exception as exc: + logger.error("OpenResponses request failed: %s", exc) + return JSONResponse( + {"error": {"message": str(exc), "type": "server_error"}}, + status_code=500, + headers={SESSION_KEY_HEADER: session_key}, + ) + + @staticmethod + def _extract_message(input_items: list[dict[str, Any]], instructions: str) -> str: + parts: list[str] = [] + if instructions: + parts.append(instructions) + for item in input_items: + if item.get("type") == "message": + content = item.get("content", "") + if isinstance(content, list): + for block in content: + if ( + isinstance(block, dict) + and block.get("type") == "input_text" + ): + parts.append(block.get("text", "")) + elif isinstance(content, str): + parts.append(content) + elif item.get("type") == "function_call_output": + parts.append(f"[tool result] {item.get('output', '')}") + return "\n".join(parts) + + @staticmethod + def _build_output_items(result: dict[str, Any]) -> list[dict[str, Any]]: + """Translate a structured turn result into OpenAI Responses output items. + + A non-empty ``summary`` becomes an assistant ``message`` item; each + ``tool_call`` event becomes a ``function_call`` item. + """ + output_items: list[dict[str, Any]] = [] + summary = result.get("summary", "") + if summary: + output_items.append( + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": summary}], + } + ) + for evt in result.get("events", []): + if evt.get("type") == "tool_call": + output_items.append( + { + "type": "function_call", + "name": evt.get("name", ""), + "arguments": evt.get("args", ""), + } + ) + return output_items + + @staticmethod + async def _responses_sse( + response_id: str, + model: str, + output_items: list[dict[str, Any]], + metadata: dict[str, Any], + ) -> AsyncIterator[bytes]: + """Emit collected output items as an OpenAI Responses-format SSE stream. + + This is a re-encoding shim, not incremental generation: the structured + turn already ran to completion, so the whole text is sent in one + ``response.output_text.delta``. The terminal ``response.completed`` + event carries the same full response object the non-streaming path + returns, so a client can rely on either. + """ + + def _event(payload: dict[str, Any]) -> bytes: + return f"data: {json.dumps(payload)}\n\n".encode() + + yield _event( + { + "type": "response.created", + "response": { + "id": response_id, + "object": "response", + "status": "in_progress", + "model": model, + }, + } + ) + for item in output_items: + if item.get("type") != "message": + continue + for block in item.get("content", []): + text = block.get("text", "") + if not text: + continue + yield _event({"type": "response.output_text.delta", "delta": text}) + yield _event({"type": "response.output_text.done", "text": text}) + yield _event( + { + "type": "response.completed", + "response": { + "id": response_id, + "object": "response", + "status": "completed", + "model": model, + "output": output_items, + "metadata": metadata, + }, + } + ) + + @staticmethod + def _derive_session_key(user: str, model: str) -> str: + if user: + return f"agent:{model or 'default'}:{user}" + return f"agent:{model or 'default'}:{uuid.uuid4().hex[:8]}" + + +def mount_bridge( + app: FastAPI, + bridge: OpenResponsesBridge, + admin_api_key: str = DEFAULT_ADMIN_API_KEY, +) -> None: + auth = make_admin_dependency(admin_api_key) + + @app.post("/v1/responses", dependencies=[Depends(auth)]) + async def responses_endpoint(request: Request): + return await bridge.handle_request(request) + + @app.on_event("shutdown") + async def shutdown_bridge(): + await bridge.close() + + +# Metadata key under which the chat bridge stashes the full, original +# ``/v1/chat/completions`` request body. Raw-passthrough agents (those +# implementing ``AgentRunnable.stream``) read it from ``request.metadata`` and +# forward it verbatim to their upstream, so no field (``ext_info`` etc.) is lost +# in the chat → turn translation. +CHAT_REQUEST_METADATA_KEY = "chat_request" + + +class ChatCompletionsBridge(AgentBridge): + """OpenAI-compatible ``/v1/chat/completions`` → DataProxy raw stream. + + Translates a chat-completions request into a DataProxy ``turn`` and relays + the worker/agent response **byte-for-byte** (typically SSE), so any + OpenAI-compatible upstream can call the gateway exactly as it would call the + backing agent directly — no client change, exact wire format preserved. The + DataProxy uses the response ``Content-Type`` to tell a raw stream apart from + a structured turn, so both protocols share the one ``turn`` endpoint. + + The full original request body is forwarded inside the turn ``metadata`` + (see :data:`CHAT_REQUEST_METADATA_KEY`) for the agent to replay verbatim. + + **Session model.** ``/v1/chat/completions`` is a stateless protocol: the + client carries the full ``messages`` history on every request, so the + framework does not store conversation history for this path. What the + framework *does* provide is **route affinity** — a stable ``session_key`` + makes the Router pin all of that session's requests to the same + DataProxy/Worker, letting the Worker-side agent reuse its per-session state + (sandbox, agent instance, KV cache). The key is resolved per request by + :meth:`_resolve_session_key`: + + 1. explicit ``X-AReaL-Session-Key`` request header, else + 2. derived from the OpenAI ``user`` field (``chat:{model}:{user}``). + + One of the two is **required** (mirroring :class:`OpenResponsesBridge`); a + request carrying neither is rejected with ``400`` rather than silently + minted a random key, so route affinity is always explicit and a multi-turn + caller can never be split across Workers by accident. + + The resolved key is always echoed back on the ``X-AReaL-Session-Key`` + response header so the caller can pin subsequent turns to it. + """ + + def __init__( + self, router_addr: str, admin_api_key: str = DEFAULT_ADMIN_API_KEY + ) -> None: + self._router_addr = router_addr + self._auth_headers = admin_headers(admin_api_key) + self._http = httpx.AsyncClient(timeout=600.0) + + async def close(self) -> None: + await self._http.aclose() + + async def handle_request(self, request: Request) -> Any: + try: + body = await request.json() + except Exception: + return JSONResponse( + {"error": {"message": "invalid JSON body", "type": "invalid_request"}}, + status_code=400, + ) + + messages = body.get("messages", []) + session_key = self._resolve_session_key(request, body) + if session_key is None: + return JSONResponse( + { + "error": { + "message": ( + "either the 'X-AReaL-Session-Key' header or the " + "'user' field is required for session affinity" + ), + "type": "invalid_request", + } + }, + status_code=400, + ) + run_id = generate_run_id() + message = self._extract_last_user_text(messages) + metadata = {CHAT_REQUEST_METADATA_KEY: body} + + try: + route_resp = await self._http.post( + f"{self._router_addr}/route", + json={"session_key": session_key}, + headers=self._auth_headers, + ) + route_resp.raise_for_status() + data_proxy_addr = route_resp.json()["data_proxy_addr"] + except Exception as exc: + logger.error("ChatCompletions routing failed: %s", exc) + return JSONResponse( + {"error": {"message": str(exc), "type": "server_error"}}, + status_code=502, + headers={SESSION_KEY_HEADER: session_key}, + ) + + # Open the DataProxy stream manually so its status/headers are known + # before we build the StreamingResponse, then relay the body verbatim. + turn_body: dict[str, Any] = { + "message": message, + "run_id": run_id, + "queue_mode": "collect", + "metadata": metadata, + } + # Opt-in self-evolution (same contract as OpenResponsesBridge): forward + # the caller-supplied inference-routing fields when present so the + # agent's LLM calls flow through AReaL's inference service under the + # caller's own ``sk-sess-*``. The DataProxy opts the turn in by the + # presence of these fields; absent them the stream is forwarded + # byte-for-byte exactly as before. + for key in ( + "inf_base_url", + "inf_model", + "session_api_key", + ): + if key in body: + turn_body[key] = body[key] + + req = self._http.build_request( + "POST", + f"{data_proxy_addr}/session/{session_key}/turn", + json=turn_body, + ) + try: + resp = await self._http.send(req, stream=True) + except Exception as exc: + logger.error("ChatCompletions upstream stream failed: %s", exc) + return JSONResponse( + {"error": {"message": str(exc), "type": "server_error"}}, + status_code=502, + headers={SESSION_KEY_HEADER: session_key}, + ) + + async def _relay(): + try: + async for chunk in resp.aiter_raw(): + yield chunk + finally: + await resp.aclose() + + headers = { + k: v + for k, v in resp.headers.items() + if k.lower() not in ("content-length", "transfer-encoding", "connection") + } + headers[SESSION_KEY_HEADER] = session_key + return StreamingResponse( + _relay(), + status_code=resp.status_code, + headers=headers, + media_type=resp.headers.get("content-type"), + ) + + @staticmethod + def _resolve_session_key(request: Request, body: dict[str, Any]) -> str | None: + """Resolve the session key, in priority order (see class docstring). + + Explicit ``X-AReaL-Session-Key`` header wins so a client can pin a + conversation onto one Worker; otherwise the OpenAI ``user`` field is + derived into a key (mirroring :class:`OpenResponsesBridge`). With + neither, return ``None`` so the caller can reject the request with + ``400`` instead of minting an implicit random key. + """ + explicit = request.headers.get(SESSION_KEY_HEADER) + if explicit: + return explicit + user = body.get("user", "") + if user: + model = body.get("model") or "default" + return f"chat:{model}:{user}" + return None + + @staticmethod + def _extract_last_user_text(messages: list[dict[str, Any]]) -> str: + for item in reversed(messages): + if item.get("role") != "user": + continue + content = item.get("content", "") + if isinstance(content, str): + return content + if isinstance(content, list): + parts = [ + block.get("text", "") + for block in content + if isinstance(block, dict) and block.get("type") == "text" + ] + return "".join(parts) + return "" + + +def mount_chat_bridge(app: FastAPI, bridge: ChatCompletionsBridge) -> None: + """Mount the chat-completions bridge. + + Deliberately **not** admin-gated: upstreams call ``/v1/chat/completions`` + exactly as they call the backing agent today (no AReaL admin key). The + internal gateway → router ``/route`` hop still carries the admin header. + """ + + @app.post("/v1/chat/completions") + async def chat_completions_endpoint(request: Request): + return await bridge.handle_request(request) + + @app.on_event("shutdown") + async def shutdown_chat_bridge(): + await bridge.close() diff --git a/areal/experimental/agent_service/gateway/config.py b/areal/v2/agent_service/gateway/config.py similarity index 100% rename from areal/experimental/agent_service/gateway/config.py rename to areal/v2/agent_service/gateway/config.py diff --git a/areal/experimental/agent_service/guard/__init__.py b/areal/v2/agent_service/guard/__init__.py similarity index 76% rename from areal/experimental/agent_service/guard/__init__.py rename to areal/v2/agent_service/guard/__init__.py index 57f50162ec..6e266647ad 100644 --- a/areal/experimental/agent_service/guard/__init__.py +++ b/areal/v2/agent_service/guard/__init__.py @@ -4,11 +4,11 @@ Pure pass-through to ``areal.infra.rpc.guard``. All orchestration logic (launching Router, Gateway, Worker+DataProxy pairs) lives in the -:mod:`~areal.experimental.agent_service.controller` module. +:mod:`~areal.v2.agent_service.controller` module. Quick start:: - python -m areal.experimental.agent_service.guard \\ + python -m areal.v2.agent_service.guard \\ --experiment-name demo --trial-name run0 \\ --role agent-guard --worker-index 0 """ diff --git a/areal/experimental/agent_service/guard/__main__.py b/areal/v2/agent_service/guard/__main__.py similarity index 80% rename from areal/experimental/agent_service/guard/__main__.py rename to areal/v2/agent_service/guard/__main__.py index d311f6023b..621b29722d 100644 --- a/areal/experimental/agent_service/guard/__main__.py +++ b/areal/v2/agent_service/guard/__main__.py @@ -1,18 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 -"""CLI entrypoint: ``python -m areal.experimental.agent_service.guard``""" +"""CLI entrypoint: ``python -m areal.v2.agent_service.guard``""" from __future__ import annotations -from areal.experimental.agent_service.guard.app import ( - _state, - app, -) from areal.infra.rpc.guard.app import ( configure_state_from_args, make_base_parser, run_server, ) +from areal.v2.agent_service.guard.app import ( + _state, + app, +) def main(): diff --git a/areal/experimental/agent_service/guard/app.py b/areal/v2/agent_service/guard/app.py similarity index 91% rename from areal/experimental/agent_service/guard/app.py rename to areal/v2/agent_service/guard/app.py index b137feef91..4241dea6f4 100644 --- a/areal/experimental/agent_service/guard/app.py +++ b/areal/v2/agent_service/guard/app.py @@ -5,7 +5,7 @@ All core guard functionality (port allocation, process forking, health checks, cleanup) is provided by ``areal.infra.rpc.guard``. This module creates and exposes the Flask app and shared state instance, following -the same pattern as ``areal.experimental.inference_service.guard``. +the same pattern as ``areal.v2.inference_service.guard``. """ from __future__ import annotations diff --git a/areal/experimental/agent_service/protocol.py b/areal/v2/agent_service/protocol.py similarity index 95% rename from areal/experimental/agent_service/protocol.py rename to areal/v2/agent_service/protocol.py index 47fb7a91e7..7668131987 100644 --- a/areal/experimental/agent_service/protocol.py +++ b/areal/v2/agent_service/protocol.py @@ -18,6 +18,13 @@ from enum import Enum from typing import Any +# HTTP response header the Worker sets on a raw-passthrough turn (an agent that +# returns a :class:`~areal.v2.agent_service.types.StreamResponse`). +# The DataProxy keys its relay-vs-parse decision on this marker rather than on +# ``Content-Type`` — a *non-streaming* passthrough body is itself +# ``application/json`` and would otherwise be mistaken for a structured turn. +PASSTHROUGH_HEADER = "x-areal-passthrough" + class FrameType(str, Enum): """WebSocket frame type discriminator.""" diff --git a/areal/experimental/agent_service/router/__init__.py b/areal/v2/agent_service/router/__init__.py similarity index 100% rename from areal/experimental/agent_service/router/__init__.py rename to areal/v2/agent_service/router/__init__.py diff --git a/areal/experimental/agent_service/router/__main__.py b/areal/v2/agent_service/router/__main__.py similarity index 95% rename from areal/experimental/agent_service/router/__main__.py rename to areal/v2/agent_service/router/__main__.py index ef5e97e7a5..a5a730b504 100644 --- a/areal/experimental/agent_service/router/__main__.py +++ b/areal/v2/agent_service/router/__main__.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 -"""``python -m areal.experimental.agent_service.router``""" +"""``python -m areal.v2.agent_service.router``""" import argparse diff --git a/areal/experimental/agent_service/router/app.py b/areal/v2/agent_service/router/app.py similarity index 100% rename from areal/experimental/agent_service/router/app.py rename to areal/v2/agent_service/router/app.py diff --git a/areal/experimental/agent_service/router/client.py b/areal/v2/agent_service/router/client.py similarity index 100% rename from areal/experimental/agent_service/router/client.py rename to areal/v2/agent_service/router/client.py diff --git a/areal/experimental/agent_service/router/config.py b/areal/v2/agent_service/router/config.py similarity index 100% rename from areal/experimental/agent_service/router/config.py rename to areal/v2/agent_service/router/config.py diff --git a/areal/v2/agent_service/types.py b/areal/v2/agent_service/types.py new file mode 100644 index 0000000000..e45f7a8199 --- /dev/null +++ b/areal/v2/agent_service/types.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Public types for the Agent Service protocol.""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +from dataclasses import dataclass, field +from typing import Any, Protocol, runtime_checkable + +from .protocol import QueueMode + + +@dataclass +class AgentRequest: + """Structured request passed to the agent. + + Core fields are stable protocol-level attributes. Framework-specific + parameters should go in *metadata*. + + Reserved metadata keys: + ``areal_inference``: present when the turn opts into AReaL's own + inference service for self-evolution (the turn carries the + inference-routing fields ``inf_base_url`` + ``session_api_key``). + Value is ``{"base_url", "api_key", "model"}`` where ``api_key`` is + the per-session ``sk-sess-*`` the **caller** obtained itself (e.g. + via its own ``/rl/start_session``) and passed in on the request. The + Agent Service does not talk to the training side; it merely + forwards these fields through. Agents should route their internal + LLM calls to this upstream so the trajectory's tokens/logprobs are + captured for training. + ``chat_request``: present on the ``/v1/chat/completions`` path; the + full original request body, so an agent that fronts an + OpenAI-compatible upstream can replay it verbatim and return a + :class:`StreamResponse` for byte-for-byte relay. + """ + + message: str + session_key: str + run_id: str + history: list[dict[str, Any]] = field(default_factory=list) + queue_mode: QueueMode = QueueMode.COLLECT + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class AgentResponse: + """Structured result returned by the agent.""" + + summary: str = "" + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class StreamResponse: + """Raw streaming response from an agent, passed through verbatim. + + The structured channel (``run`` + :class:`EventEmitter` → ``AgentResponse``) + parses an agent's output into deltas/tool-calls. Some agents instead expose + an OpenAI-compatible upstream whose response (often SSE) must reach the + caller **byte-for-byte** — re-encoding it through the structured channel + would drop fields (tool_calls, finish_reason, usage, ...) and break clients + that expect the exact wire format. + + An agent opts into this behaviour by returning a :class:`StreamResponse` + (instead of an :class:`AgentResponse`) from :meth:`AgentRunnable.run`. The + Worker and DataProxy relay ``status_code`` / ``headers`` / ``body`` through + the single ``/run`` and ``/session/{key}/turn`` endpoints without inspecting + the payload — the Worker tags the response with the ``x-areal-passthrough`` + marker header so they relay it verbatim (structured turns are parsed). The + marker, not ``Content-Type``, drives the decision, so a *non-streaming* + passthrough whose body is ``application/json`` is still relayed byte-for-byte + rather than mistaken for a structured turn. + """ + + status_code: int + headers: dict[str, str] + body: AsyncIterator[bytes] + + +class EventEmitter(Protocol): + """Callback interface for streaming events from agent to caller.""" + + async def emit_delta(self, text: str) -> None: ... + async def emit_tool_call(self, name: str, args: str) -> None: ... + async def emit_tool_result(self, name: str, result: str) -> None: ... + + +@runtime_checkable +class AgentRunnable(Protocol): + """Minimal protocol for pluggable agent implementations. + + Agent classes are loaded via + :func:`~areal.utils.dynamic_import.import_from_string` at worker startup. + The framework handles its own tool execution, memory, and LLM + interaction — the Agent Service only provides session lifecycle and + event streaming. + + ``run`` is the single entry point and may return **either** shape, + chosen per turn by the agent: + + - :class:`AgentResponse` — the structured channel. The agent reports + incremental output through the ``emitter`` and returns a final + summary/metadata; the Worker serialises it to JSON and the DataProxy + rebuilds conversation history from the emitted events. This backs the + ``/v1/responses`` (and WebSocket) protocol. + - :class:`StreamResponse` — the raw-passthrough channel. The agent fronts + an OpenAI-compatible upstream whose response must reach the caller + byte-for-byte (e.g. SSE chat completions); it returns the upstream's + ``status_code`` / ``headers`` / ``body`` and the Worker / DataProxy relay + them verbatim. This backs the ``/v1/chat/completions`` protocol. + + The agent decides which to return from the request (e.g. the presence of + ``metadata['chat_request']``, or a ``stream`` flag in the original body), + so a single ``run`` implementation can serve every protocol and both + streaming and non-streaming modes. + + The following methods are optional and discovered via ``getattr`` at + runtime — implement them to participate in training-related lifecycle: + + - ``async close_session(session_key)`` — release per-session state + when a session is closed by the DataProxy. + - ``async close_all_sessions()`` — clean up everything on worker + shutdown. + """ + + async def run( + self, + request: AgentRequest, + *, + emitter: EventEmitter, + ) -> AgentResponse | StreamResponse: ... diff --git a/areal/experimental/agent_service/worker/__init__.py b/areal/v2/agent_service/worker/__init__.py similarity index 100% rename from areal/experimental/agent_service/worker/__init__.py rename to areal/v2/agent_service/worker/__init__.py diff --git a/areal/experimental/agent_service/worker/__main__.py b/areal/v2/agent_service/worker/__main__.py similarity index 89% rename from areal/experimental/agent_service/worker/__main__.py rename to areal/v2/agent_service/worker/__main__.py index f1cc525bdf..da840adb6f 100644 --- a/areal/experimental/agent_service/worker/__main__.py +++ b/areal/v2/agent_service/worker/__main__.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 -"""``python -m areal.experimental.agent_service.worker`` +"""``python -m areal.v2.agent_service.worker`` Start a standalone Agent Worker process. The Controller forks this via Guard to create Worker+DataProxy pairs. - python -m areal.experimental.agent_service.worker \ + python -m areal.v2.agent_service.worker \ --agent examples.agent_service.agent.ClaudeAgent \ --host 127.0.0.1 --port 9000 """ diff --git a/areal/experimental/agent_service/worker/app.py b/areal/v2/agent_service/worker/app.py similarity index 58% rename from areal/experimental/agent_service/worker/app.py rename to areal/v2/agent_service/worker/app.py index 55507bc558..6cb6881ccb 100644 --- a/areal/experimental/agent_service/worker/app.py +++ b/areal/v2/agent_service/worker/app.py @@ -8,13 +8,18 @@ from typing import Any from fastapi import FastAPI -from fastapi.responses import JSONResponse +from fastapi.responses import JSONResponse, StreamingResponse from areal.utils import logging from areal.utils.dynamic_import import import_from_string -from ..protocol import QueueMode -from ..types import AgentRequest, AgentResponse, AgentRunnable +from ..protocol import PASSTHROUGH_HEADER, QueueMode +from ..types import ( + AgentRequest, + AgentResponse, + AgentRunnable, + StreamResponse, +) logger = logging.getLogger("AgentWorker") @@ -67,6 +72,20 @@ async def shutdown(): @app.post("/run") async def run(body: dict[str, Any]): + """Single agent entry point for every protocol and streaming mode. + + Calls the agent's ``run`` and relays whichever shape it returns: + + - :class:`StreamResponse` — raw passthrough; ``status_code`` / + ``headers`` / ``body`` are forwarded untouched so the caller gets the + upstream's exact wire format (e.g. SSE chat completions). The + response carries the :data:`PASSTHROUGH_HEADER` marker so the + DataProxy relays it verbatim instead of parsing it — this works even + for a *non-streaming* passthrough whose body is ``application/json``. + - :class:`AgentResponse` — structured JSON ``{summary, metadata, + events}`` (``application/json``); the DataProxy rebuilds history from + ``events``. + """ request = AgentRequest( message=body.get("message", ""), session_key=body.get("session_key", ""), @@ -79,7 +98,9 @@ async def run(body: dict[str, Any]): emitter = _CollectingEmitter() try: - response: AgentResponse = await agent.run(request, emitter=emitter) + response: AgentResponse | StreamResponse = await agent.run( + request, emitter=emitter + ) except Exception as exc: logger.exception("Agent run failed (session=%s)", request.session_key) return JSONResponse( @@ -87,6 +108,25 @@ async def run(body: dict[str, Any]): status_code=500, ) + if isinstance(response, StreamResponse): + # Drop hop-by-hop / length headers that would conflict with chunked + # relaying; FastAPI/uvicorn set framing headers themselves. + headers = { + k: v + for k, v in response.headers.items() + if k.lower() + not in ("content-length", "transfer-encoding", "connection") + } + # Mark the turn as raw-passthrough so the DataProxy relays it + # verbatim regardless of its Content-Type. + headers[PASSTHROUGH_HEADER] = "1" + return StreamingResponse( + response.body, + status_code=response.status_code, + headers=headers, + media_type=response.headers.get("content-type"), + ) + return {**asdict(response), "events": emitter.events} return app diff --git a/areal/experimental/agent_service/worker/config.py b/areal/v2/agent_service/worker/config.py similarity index 100% rename from areal/experimental/agent_service/worker/config.py rename to areal/v2/agent_service/worker/config.py diff --git a/areal/experimental/inference_service/gateway/__init__.py b/areal/v2/cli/__init__.py similarity index 100% rename from areal/experimental/inference_service/gateway/__init__.py rename to areal/v2/cli/__init__.py diff --git a/areal/v2/cli/__main__.py b/areal/v2/cli/__main__.py new file mode 100644 index 0000000000..02f5fa0c1f --- /dev/null +++ b/areal/v2/cli/__main__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 + +from areal.v2.cli.cli import cli + +if __name__ == "__main__": + cli() diff --git a/areal/v2/cli/agent/__init__.py b/areal/v2/cli/agent/__init__.py new file mode 100644 index 0000000000..724728eb7d --- /dev/null +++ b/areal/v2/cli/agent/__init__.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path + +import click + +from areal.v2.cli.agent.commands.ps import ps_cmd +from areal.v2.cli.agent.commands.run import run_cmd +from areal.v2.cli.agent.commands.status import status_cmd +from areal.v2.cli.agent.commands.stop import stop_cmd +from areal.v2.cli.agent.config import load_click_default_map +from areal.v2.cli.agent.lifecycle import agent_lifecycle +from areal.v2.cli.commands.logs import LogsCommand + + +@click.group(help="Manage agent services.") +@click.option( + "--config", + "config_file", + type=click.Path(exists=True, dir_okay=False, path_type=Path), + default=None, + help="Extra TOML file merged on top of ~/.areal/agent/config.toml.", +) +@click.pass_context +def agent(ctx: click.Context, config_file: Path | None) -> None: + ctx.default_map = load_click_default_map(extra=config_file) + + +agent.add_command(run_cmd) +agent.add_command(stop_cmd) +agent.add_command(status_cmd) +agent.add_command(ps_cmd) +agent.add_command(LogsCommand(lifecycle=agent_lifecycle).build()) diff --git a/areal/v2/cli/agent/cli_guide.md b/areal/v2/cli/agent/cli_guide.md new file mode 100644 index 0000000000..1524d6f698 --- /dev/null +++ b/areal/v2/cli/agent/cli_guide.md @@ -0,0 +1,180 @@ +# AReaL Agent Service CLI + +`areal agent` is the agent subcommand group under the top-level `areal` CLI. It launches +a set of agent service processes on the local machine (gateway / router + N +worker/data-proxy pairs) so that an upstream application can interact with the agent +over HTTP. Its shape is very similar to `areal inf`, but it serves an agent (multi-turn +interaction with session state) rather than stateless inference. + +## Basic concepts + +A running agent service has the following components: + +- `gateway`: exposes session and agent APIs to the outside. One per service. +- `router`: routes requests to a data-proxy based on load. One per service. +- `worker[i]`: the process that actually runs user agent code. The agent class is + imported from the `module.path` form passed via `--agent`. +- `data-proxy[i]`: the session-management / accounting layer in front of `worker[i]`. + Each worker is paired with one proxy to form a **pair**.`--num-pairs` controls the + number of replicas. + +Request flow: + +``` +client → gateway → router → data-proxy[i] → worker[i] → agent code + └ session lifecycle / record ┘ +``` + +The CLI's local state is stored under `~/.areal/agent/` by default. The root directory +can be overridden via `AREAL_HOME`: + +```bash +export AREAL_HOME=/path/to/areal-home +``` + +## Launching the service + +Minimum launch — a single (worker, proxy) pair: + +```bash +areal agent run \ + --service default \ + --agent my_package.my_agent.MyAgent \ + --num-pairs 1 \ + --admin-api-key areal-agent-admin +``` + +`--agent` is required; it is the import path the worker process uses to load the agent +class. + +Force-start by clearing stale state: + +```bash +areal agent run --service default --agent ... --force +``` + +## Inspecting service state + +List all agent services on the local machine: + +```bash +areal agent ps +areal agent ps --all # include stale rows +areal agent ps --json +``` + +Output columns: `SERVICE / STATUS / GATEWAY / AGENT`. + +Drill into a single service for per-component health: + +```bash +areal agent status --service default +``` + +The output includes the gateway, router, and each pair's worker + proxy. `--watch` mode +refreshes on an interval (default 2 seconds): + +```bash +areal agent status --service default --watch --interval 1 +``` + +JSON mode plays well with jq: + +```bash +areal agent status --service default --json | jq '.pairs[].worker' +``` + +## Talking to the service + +The CLI **does not** manage how an application talks to the service —applications hit +the gateway HTTP endpoints directly. The status command tells you the gateway URL: + +```bash +GATEWAY_URL=$(areal agent status --service default --json | jq -r '.gateway.url') +echo "gateway at $GATEWAY_URL" +``` + +The application then sends requests against that URL with `--admin-api-key` (or a +session key obtained from the gateway). + +## Logs + +Each component writes a separate log file: + +```bash +areal agent logs --service default --component gateway -f +areal agent logs --service default --component router -f +areal agent logs --service default --component worker-0 -f +areal agent logs --service default --component proxy-0 -f +``` + +Naming convention: + +- `gateway` / `router`: service-level singletons +- `worker-` / `proxy-`: the worker / data-proxy of the i-th pair (istarts from 0) + +If `--component` is wrong, the CLI prints the available names. `-f` uses `tail -F` +semantics. + +## Stopping + +```bash +areal agent stop --service default +``` + +The default is a two-phase shutdown: SIGTERM, wait `--grace-period`(10s), then SIGKILL. +Immediate SIGKILL: + +```bash +areal agent stop --service default --force +``` + +`--keep-state` preserves the state file (kills processes but leaves the on-disk +`.json` alone): + +```bash +areal agent stop --service default --keep-state +``` + +## Configuration file + +`areal agent` reads `~/.areal/agent/config.toml` as defaults on startup;additional +config files can be passed in: + +```bash +areal agent --config ./my-agent.toml run --service default --agent ... +``` + +Example: + +```toml +[default] +service = "default" +admin_api_key = "areal-agent-admin" +log_level = "info" + +[run] +agent = "my_package.my_agent.MyAgent" +num_pairs = 2 +setup_timeout = 120 +health_poll_interval = 5 +drain_timeout = 30 +session_timeout = 1800 +``` + +Precedence: \*\*CLI flag > TOML passed via `--config` > `~/.areal/agent/config.toml` + +> hardcoded defaults\*\*. + +## Not implemented yet + +The current `areal agent` does **not** include: + +- Session-level CLI operations (start session / set reward / export trajectory) — these + are tightly coupled with the application and are handled by the application talking + directly to the gateway HTTP +- Automatic failure recovery / heartbeat monitoring — `status` is on-demand; it does not + continuously observe component health. If a worker dies, users have to discover it by + running `status` or checking logs. +- Distributed scheduling — only local processes on the local machine; k8s / slurm and + friends are out of scope for the current CLI. diff --git a/areal/v2/cli/agent/client.py b/areal/v2/cli/agent/client.py new file mode 100644 index 0000000000..da22c3a04b --- /dev/null +++ b/areal/v2/cli/agent/client.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import Any + +from areal.v2.cli.client import BaseHTTPClient + + +class RouterClient(BaseHTTPClient): + def register_proxy(self, addr: str, *, timeout: float = 10.0) -> dict[str, Any]: + return self._post("/register", payload={"addr": addr}, timeout=timeout) diff --git a/areal/experimental/inference_service/guard/__init__.py b/areal/v2/cli/agent/commands/__init__.py similarity index 100% rename from areal/experimental/inference_service/guard/__init__.py rename to areal/v2/cli/agent/commands/__init__.py diff --git a/areal/v2/cli/agent/commands/ps.py b/areal/v2/cli/agent/commands/ps.py new file mode 100644 index 0000000000..b2efe73f2b --- /dev/null +++ b/areal/v2/cli/agent/commands/ps.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import click + +from areal.v2.cli.agent.lifecycle import agent_lifecycle +from areal.v2.cli.process import pid_alive +from areal.v2.cli.utils import json_or_table + + +@click.command(name="ps", help="List locally known agent services.") +@click.option("--json", "as_json", is_flag=True) +@click.option("--all", "include_all", is_flag=True, help="Include stale services.") +def ps_cmd(as_json: bool, include_all: bool) -> None: + raise SystemExit(do_ps(as_json=as_json, include_all=include_all) or 0) + + +def do_ps(*, as_json: bool, include_all: bool) -> int: + rows: list[dict] = [] + for service in agent_lifecycle.list_services(): + try: + state = agent_lifecycle.load_state(service) + except Exception: + if include_all: + rows.append({"service": service, "status": "stale"}) + continue + running = pid_alive(state.gateway.pid) + if running or include_all: + rows.append( + { + "service": service, + "status": "running" if running else "stale", + "gateway_url": state.gateway.url, + "agent": state.agent, + } + ) + + json_or_table(rows, as_json=as_json, table_renderer=_print_table) + return 0 + + +def _print_table(rows: list[dict]) -> None: + if not rows: + click.echo("no agent services") + return + cols = ("SERVICE", "STATUS", "GATEWAY", "AGENT") + table = [ + ( + row["service"], + row["status"], + row.get("gateway_url", ""), + row.get("agent", ""), + ) + for row in rows + ] + widths = [max(len(str(r[i])) for r in (cols, *table)) for i in range(4)] + fmt = " ".join(f"{{:<{w}}}" for w in widths) + click.echo(fmt.format(*cols)) + for row in table: + click.echo(fmt.format(*row)) diff --git a/areal/v2/cli/agent/commands/run.py b/areal/v2/cli/agent/commands/run.py new file mode 100644 index 0000000000..7e8ca6628f --- /dev/null +++ b/areal/v2/cli/agent/commands/run.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import click + +from areal.v2.cli.agent.launcher import launch_agent_stack +from areal.v2.cli.agent.lifecycle import agent_lifecycle +from areal.v2.cli.agent.state import ServiceState +from areal.v2.cli.client import ServiceHTTPError, ServiceUnreachable +from areal.v2.cli.process import kill_pids +from areal.v2.cli.state import DEFAULT_SERVICE +from areal.v2.cli.utils import register_cli_logger + +logger = register_cli_logger("AgentCli") + + +@click.command(name="run", help="Launch an agent service.") +@click.option("--service", default=DEFAULT_SERVICE, show_default=True) +@click.option("--agent", default=None, help="Agent import path.") +@click.option("--num-pairs", type=int, default=1, show_default=True) +@click.option("--admin-api-key", default="areal-agent-admin", show_default=True) +@click.option("--setup-timeout", type=float, default=120.0, show_default=True) +@click.option("--health-poll-interval", type=float, default=5.0, show_default=True) +@click.option("--drain-timeout", type=float, default=30.0, show_default=True) +@click.option("--session-timeout", type=float, default=1800.0, show_default=True) +@click.option( + "--log-level", + type=click.Choice(["debug", "info", "warning", "error"]), + default="info", + show_default=True, +) +@click.option("--force", is_flag=True, help="Replace stale or running service state.") +def run_cmd(**opts) -> None: + raise SystemExit(do_run(**opts) or 0) + + +def do_run( + *, + service: str, + agent: str | None, + num_pairs: int, + admin_api_key: str, + setup_timeout: float, + health_poll_interval: float, + drain_timeout: float, + session_timeout: float, + log_level: str, + force: bool, +) -> int: + if not agent: + raise click.UsageError("--agent is required") + + if force: + agent_lifecycle.force_replace_slot(service, grace_s=5.0) + else: + agent_lifecycle.refuse_if_running(service) + + launched: ServiceState | None = None + try: + launched = launch_agent_stack( + service=service, + agent=agent, + admin_api_key=admin_api_key, + num_pairs=num_pairs, + setup_timeout=setup_timeout, + session_timeout=session_timeout, + health_poll_interval=health_poll_interval, + drain_timeout=drain_timeout, + log_level=log_level, + ) + launched.save() + except (ServiceHTTPError, ServiceUnreachable, RuntimeError, ValueError) as exc: + if launched is not None: + kill_pids( + [pid for _, handle in launched.components() if (pid := handle.pid) > 0], + grace_s=5.0, + ) + raise click.ClickException(f"failed to launch agent service: {exc}") from exc + + logger.info("service=%s gateway=%s", service, launched.gateway.url) + return 0 diff --git a/areal/v2/cli/agent/commands/status.py b/areal/v2/cli/agent/commands/status.py new file mode 100644 index 0000000000..60b5ad596b --- /dev/null +++ b/areal/v2/cli/agent/commands/status.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import time + +import click + +from areal.v2.cli.agent.lifecycle import agent_lifecycle +from areal.v2.cli.process import pid_alive +from areal.v2.cli.state import SupportsComponentProbe +from areal.v2.cli.status import ColumnSpec, StatusReporter +from areal.v2.cli.utils import json_or_table + + +@click.command(name="status", help="Show agent service health.") +@click.option("--service", default=None) +@click.option("--watch", is_flag=True) +@click.option("--interval", type=float, default=2.0, show_default=True) +@click.option("--json", "as_json", is_flag=True) +def status_cmd( + service: str | None, watch: bool, interval: float, as_json: bool +) -> None: + raise SystemExit( + do_status(service=service, watch=watch, interval=interval, as_json=as_json) or 0 + ) + + +def do_status( + *, service: str | None, watch: bool, interval: float, as_json: bool +) -> int: + name = agent_lifecycle.resolve_service_name(service) + while True: + _emit_once(name, as_json=as_json) + if not watch: + return 0 + time.sleep(interval) + + +def _emit_once(service: str, *, as_json: bool) -> None: + path = agent_lifecycle.state_path(service) + if not path.exists(): + payload = {"service": service, "running": False, "components": []} + json_or_table( + payload, + as_json=as_json, + table_renderer=lambda p: click.echo( + f"service {p['service']!r} is not running" + ), + ) + return + try: + state = agent_lifecycle.load_state(service) + except Exception as exc: + payload = { + "service": service, + "running": False, + "error": f"failed to read state: {exc}", + "components": [], + } + json_or_table( + payload, + as_json=as_json, + table_renderer=lambda p: click.echo( + f"service {p['service']!r}: {p['error']}" + ), + ) + return + + components = list(state.components()) + reporter = StatusReporter(components, _columns(service)) + alive = reporter.probe_all() + + if as_json: + payload = { + "service": service, + "running": pid_alive(state.gateway.pid), + "gateway_url": state.gateway.url, + "router_url": state.router.url, + "components": reporter.json_snapshot(alive), + } + click.echo(_json_indent(payload)) + return + reporter.print_table( + reporter.render_rows(alive), + header_line=f"service: {service} gateway: {state.gateway.url}", + ) + + +def _columns(service: str) -> list[ColumnSpec]: + def _service(label: str, _: SupportsComponentProbe, __: bool) -> str: + del label + return service + + def _component(label: str, _: SupportsComponentProbe, __: bool) -> str: + return label + + def _status(_: str, __: SupportsComponentProbe, alive: bool) -> str: + return "ok" if alive else "down" + + def _addr(_: str, handle: SupportsComponentProbe, __: bool) -> str: + return handle.addr or "-" + + def _pid(_: str, handle: SupportsComponentProbe, __: bool) -> str: + return str(handle.pid) if handle.pid > 0 else "-" + + return [ + ColumnSpec("SERVICE", _service), + ColumnSpec("COMPONENT", _component), + ColumnSpec("STATUS", _status), + ColumnSpec("ADDR", _addr), + ColumnSpec("PID", _pid), + ] + + +def _json_indent(payload: dict) -> str: + import json + + return json.dumps(payload, indent=2, default=str) diff --git a/areal/v2/cli/agent/commands/stop.py b/areal/v2/cli/agent/commands/stop.py new file mode 100644 index 0000000000..5cea6302c4 --- /dev/null +++ b/areal/v2/cli/agent/commands/stop.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import click + +from areal.v2.cli.agent.lifecycle import agent_lifecycle +from areal.v2.cli.agent.state import ServiceState +from areal.v2.cli.process import kill_pids + + +@click.command(name="stop", help="Stop an agent service.") +@click.option("--service", default=None) +@click.option("--grace-period", type=float, default=10.0, show_default=True) +@click.option("--keep-state", is_flag=True) +@click.option("--force", is_flag=True, help="SIGKILL immediately.") +def stop_cmd( + service: str | None, grace_period: float, keep_state: bool, force: bool +) -> None: + raise SystemExit( + do_stop( + service=service, + grace_period=grace_period, + keep_state=keep_state, + force=force, + ) + or 0 + ) + + +def do_stop( + *, + service: str | None, + grace_period: float, + keep_state: bool, + force: bool, +) -> int: + name = agent_lifecycle.resolve_service_name(service) + path = agent_lifecycle.state_path(name) + if not path.exists(): + click.echo(f"service {name!r} is not running") + return 0 + + try: + state = agent_lifecycle.load_state(name) + except Exception: + if not keep_state: + ServiceState.remove(name) + click.echo(f"removed stale state for {name!r}") + return 0 + + pids = [pid for _, h in state.components() if (pid := h.pid) > 0] + kill_pids(pids, grace_s=0.0 if force else grace_period) + if not keep_state: + ServiceState.remove(name) + click.echo(f"service {name!r} stopped") + return 0 diff --git a/areal/v2/cli/agent/config.py b/areal/v2/cli/agent/config.py new file mode 100644 index 0000000000..d493604ea9 --- /dev/null +++ b/areal/v2/cli/agent/config.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path + +from areal.v2.cli.agent.state import AGENT_NAMESPACE +from areal.v2.cli.config import BindingMap, ConfigLoader + +# (section, key) -> (verb_or_verbs, click_option_name). +# Anything not listed in the agent's TOML is silently ignored. +AGENT_BINDINGS: BindingMap = { + ("default", "service"): (("run", "stop", "status", "ps", "logs"), "service"), + ("default", "admin_api_key"): ("run", "admin_api_key"), + ("default", "log_level"): ("run", "log_level"), + ("run", "agent"): ("run", "agent"), + ("run", "num_pairs"): ("run", "num_pairs"), + ("run", "setup_timeout"): ("run", "setup_timeout"), + ("run", "health_poll_interval"): ("run", "health_poll_interval"), + ("run", "drain_timeout"): ("run", "drain_timeout"), + ("run", "session_timeout"): ("run", "session_timeout"), +} + + +agent_config_loader = ConfigLoader(namespace=AGENT_NAMESPACE, bindings=AGENT_BINDINGS) + + +def load_click_default_map(extra: Path | None = None) -> dict: + return agent_config_loader.load_click_default_map(extra=extra) diff --git a/areal/v2/cli/agent/launcher.py b/areal/v2/cli/agent/launcher.py new file mode 100644 index 0000000000..2ce55e2d4f --- /dev/null +++ b/areal/v2/cli/agent/launcher.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import sys +from pathlib import Path + +from areal.utils.network import find_free_ports +from areal.v2.cli.agent.client import RouterClient +from areal.v2.cli.agent.state import ( + PairState, + ProcessState, + ServiceState, + store, +) +from areal.v2.cli.process import spawn_process +from areal.v2.cli.utils import wait_http_health + + +def launch_agent_stack( + *, + service: str, + agent: str, + admin_api_key: str, + num_pairs: int, + setup_timeout: float, + session_timeout: float, + health_poll_interval: float, + drain_timeout: float, + log_level: str, +) -> ServiceState: + log_dir = store.logs_dir(service) + + # Reserve every port we need up front from the non-ephemeral pool to + # avoid TOCTOU collisions between successive bind(0) calls. + ports = find_free_ports(2 + 2 * num_pairs) + router_port = ports[0] + gateway_port = ports[1] + pair_ports = ports[2:] + + router_log = log_dir / "router.log" + router_pid = _spawn_router( + host="127.0.0.1", + port=router_port, + admin_api_key=admin_api_key, + log_level=log_level, + log_file=router_log, + ) + router_url = f"http://127.0.0.1:{router_port}" + wait_http_health(router_url, pid=router_pid, timeout=setup_timeout, label="router") + + pairs: list[PairState] = [] + router = RouterClient(router_url, admin_api_key) + for idx in range(num_pairs): + pair = _spawn_pair( + index=idx, + agent=agent, + worker_port=pair_ports[2 * idx], + proxy_port=pair_ports[2 * idx + 1], + session_timeout=session_timeout, + log_level=log_level, + log_dir=log_dir, + setup_timeout=setup_timeout, + ) + router.register_proxy(pair.data_proxy.url) + pairs.append(pair) + + gateway_log = log_dir / "gateway.log" + gateway_pid = _spawn_gateway( + host="127.0.0.1", + port=gateway_port, + router_url=router_url, + admin_api_key=admin_api_key, + log_level=log_level, + log_file=gateway_log, + ) + gateway_url = f"http://127.0.0.1:{gateway_port}" + wait_http_health( + gateway_url, pid=gateway_pid, timeout=setup_timeout, label="gateway" + ) + + return ServiceState( + service=service, + launch_mode="detached", + agent=agent, + admin_api_key=admin_api_key, + gateway=ProcessState( + component="gateway", + pid=gateway_pid, + url=gateway_url, + log_file=str(gateway_log), + ), + router=ProcessState( + component="router", + pid=router_pid, + url=router_url, + log_file=str(router_log), + ), + pairs=pairs, + session_timeout=session_timeout, + health_poll_interval=health_poll_interval, + drain_timeout=drain_timeout, + ) + + +def _spawn_router( + *, + host: str, + port: int, + admin_api_key: str, + log_level: str, + log_file: Path, +) -> int: + return spawn_process( + [ + sys.executable, + "-m", + "areal.v2.agent_service.router", + "--host", + host, + "--port", + str(port), + "--admin-api-key", + admin_api_key, + "--log-level", + log_level, + ], + log_file, + ) + + +def _spawn_gateway( + *, + host: str, + port: int, + router_url: str, + admin_api_key: str, + log_level: str, + log_file: Path, +) -> int: + return spawn_process( + [ + sys.executable, + "-m", + "areal.v2.agent_service.gateway", + "--host", + host, + "--port", + str(port), + "--router-addr", + router_url, + "--admin-api-key", + admin_api_key, + "--log-level", + log_level, + ], + log_file, + ) + + +def _spawn_pair( + *, + index: int, + agent: str, + worker_port: int, + proxy_port: int, + session_timeout: float, + log_level: str, + log_dir: Path, + setup_timeout: float, +) -> PairState: + worker_log = log_dir / f"worker-{index}.log" + worker_pid = spawn_process( + [ + sys.executable, + "-m", + "areal.v2.agent_service.worker", + "--agent", + agent, + "--host", + "127.0.0.1", + "--port", + str(worker_port), + "--log-level", + log_level, + ], + worker_log, + ) + worker_url = f"http://127.0.0.1:{worker_port}" + wait_http_health( + worker_url, pid=worker_pid, timeout=setup_timeout, label=f"worker-{index}" + ) + + proxy_log = log_dir / f"proxy-{index}.log" + proxy_pid = spawn_process( + [ + sys.executable, + "-m", + "areal.v2.agent_service.data_proxy", + "--worker-addr", + worker_url, + "--host", + "127.0.0.1", + "--port", + str(proxy_port), + "--session-timeout", + str(int(session_timeout)), + "--log-level", + log_level, + ], + proxy_log, + ) + proxy_url = f"http://127.0.0.1:{proxy_port}" + wait_http_health( + proxy_url, pid=proxy_pid, timeout=setup_timeout, label=f"proxy-{index}" + ) + + return PairState( + index=index, + worker=ProcessState( + component=f"worker-{index}", + pid=worker_pid, + url=worker_url, + log_file=str(worker_log), + ), + data_proxy=ProcessState( + component=f"proxy-{index}", + pid=proxy_pid, + url=proxy_url, + log_file=str(proxy_log), + ), + ) diff --git a/areal/v2/cli/agent/lifecycle.py b/areal/v2/cli/agent/lifecycle.py new file mode 100644 index 0000000000..cfd93b0076 --- /dev/null +++ b/areal/v2/cli/agent/lifecycle.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from areal.v2.cli.agent.state import AGENT_NAMESPACE, ServiceState +from areal.v2.cli.lifecycle import ServiceLifecycle + +agent_lifecycle = ServiceLifecycle( + namespace=AGENT_NAMESPACE, + state_class=ServiceState, + stop_command="areal agent stop", +) diff --git a/areal/v2/cli/agent/state.py b/areal/v2/cli/agent/state.py new file mode 100644 index 0000000000..42d7a73ace --- /dev/null +++ b/areal/v2/cli/agent/state.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import time +from collections.abc import Iterable +from dataclasses import asdict, dataclass, field + +from areal.v2.cli.process import pid_alive +from areal.v2.cli.state import ( + NamespacedStateStore, + ServiceStateBase, + SupportsComponentProbe, + atomic_write_json, +) + +AGENT_NAMESPACE = "agent" + +# Module-level store — every save/load/remove routes paths through it, +# tests with ``AREAL_HOME`` work because the store re-resolves on each +# call. +store = NamespacedStateStore(AGENT_NAMESPACE) + + +@dataclass +class ProcessState: + component: str + pid: int + url: str + log_file: str + + @property + def addr(self) -> str: + return self.url + + +@dataclass +class PairState: + index: int + worker: ProcessState + data_proxy: ProcessState + + +@dataclass +class ServiceState(ServiceStateBase): + service: str + launch_mode: str + agent: str + admin_api_key: str + gateway: ProcessState + router: ProcessState + pairs: list[PairState] + session_timeout: float = 1800.0 + health_poll_interval: float = 5.0 + drain_timeout: float = 30.0 + started_at: float = field(default_factory=time.time) + + def save(self) -> None: + atomic_write_json(store.service_state_path(self.service), asdict(self)) + store.set_current_service(self.service) + + @classmethod + def load(cls, service: str) -> ServiceState: + with open(store.service_state_path(service)) as f: + raw = json.load(f) + raw["gateway"] = ProcessState(**raw["gateway"]) + raw["router"] = ProcessState(**raw["router"]) + raw["pairs"] = [ + PairState( + index=pair["index"], + worker=ProcessState(**pair["worker"]), + data_proxy=ProcessState(**pair["data_proxy"]), + ) + for pair in raw.get("pairs", []) + ] + # ``created_at`` is the legacy field name; accept it for forward + # compatibility with any state files written before the rename. + if "created_at" in raw and "started_at" not in raw: + raw["started_at"] = raw.pop("created_at") + raw.pop("created_at", None) + for legacy in ("inf_addr", "inf_api_key", "inf_model"): + raw.pop(legacy, None) + return cls(**raw) + + @classmethod + def remove(cls, service: str) -> None: + path = store.service_state_path(service) + if path.exists(): + path.unlink() + store.clear_current_service(service) + + def gateway_alive(self) -> bool: + # The CLI treats "running" as "gateway process still alive". The + # gateway is the entry point — without it, the rest of the stack + # is unreachable even if some children survived. + return pid_alive(self.gateway.pid) + + def components(self) -> Iterable[tuple[str, SupportsComponentProbe]]: + yield "gateway", self.gateway + yield "router", self.router + for pair in self.pairs: + yield pair.worker.component, pair.worker + yield pair.data_proxy.component, pair.data_proxy diff --git a/areal/v2/cli/cli.py b/areal/v2/cli/cli.py new file mode 100644 index 0000000000..d40d39c121 --- /dev/null +++ b/areal/v2/cli/cli.py @@ -0,0 +1,24 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import click + +from areal.v2.cli.agent import agent +from areal.v2.cli.inference import inf +from areal.v2.cli.training import train +from areal.version import __version__ + + +@click.group( + context_settings={"help_option_names": ["-h", "--help"]}, + help="AReaL operator CLI.", +) +@click.version_option(__version__, prog_name="areal") +def cli() -> None: + pass + + +cli.add_command(agent) +cli.add_command(inf) +cli.add_command(train) diff --git a/areal/v2/cli/client.py b/areal/v2/cli/client.py new file mode 100644 index 0000000000..fac1ae5f39 --- /dev/null +++ b/areal/v2/cli/client.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""HTTP client base for subcommand CLIs. + +Provides: + +- ``ServiceHTTPError`` / ``ServiceUnreachable`` exceptions to distinguish + "server replied with 4xx/5xx" from "couldn't reach the server". +- A low-level ``request_json`` helper that builds the request, handles + optional Bearer auth, decodes JSON, and translates ``urllib`` errors + into the two exception types above. +- ``BaseHTTPClient`` — minimal base for per-component clients + (gateway / router / data-proxy / etc.). Subclasses add their own RPC + methods that call ``self._get`` / ``self._post`` / ``self._request``. +""" + +from __future__ import annotations + +import json +import urllib.error +import urllib.request +from typing import Any + + +class ServiceHTTPError(Exception): + """Non-2xx response from the service. Body is captured so callers + can inspect application-level error payloads.""" + + def __init__(self, status: int, body: str) -> None: + super().__init__(f"HTTP {status}: {body}") + self.status = status + self.body = body + + +class ServiceUnreachable(Exception): + """Network-level error reaching the service: connection refused, + DNS failure, timeout, etc.""" + + +def request_json( + url: str, + *, + method: str = "GET", + payload: dict | None = None, + bearer: str | None = None, + timeout: float = 5.0, +) -> dict[str, Any]: + """Issue an HTTP request returning parsed JSON. + + A ``payload`` dict triggers JSON encoding + ``Content-Type``. + ``bearer`` adds ``Authorization: Bearer ``. 2xx → returns the + parsed body (empty dict if body is empty). 4xx/5xx → ``ServiceHTTPError``. + Network errors → ``ServiceUnreachable``. + """ + + data = None + headers = {"Accept": "application/json"} + if payload is not None: + data = json.dumps(payload).encode() + headers["Content-Type"] = "application/json" + if bearer: + headers["Authorization"] = f"Bearer {bearer}" + req = urllib.request.Request(url, data=data, method=method, headers=headers) + try: + with urllib.request.urlopen(req, timeout=timeout) as resp: + body = resp.read().decode() + return json.loads(body) if body else {} + except urllib.error.HTTPError as exc: + raise ServiceHTTPError(exc.code, exc.read().decode(errors="replace")) from exc + except (urllib.error.URLError, ConnectionError, TimeoutError, OSError) as exc: + raise ServiceUnreachable(str(exc)) from exc + + +class BaseHTTPClient: + """Minimal HTTP client base. Subclasses add business RPC methods. + + Holds ``base_url`` and an optional ``api_key`` used as Bearer auth on + every helper call. ``health()`` is provided as it's universal across + all daemon types. + """ + + def __init__(self, base_url: str, api_key: str | None = None) -> None: + self.base = base_url.rstrip("/") + self.api_key = api_key + + def health(self, *, timeout: float = 2.0) -> dict[str, Any]: + return self._get("/health", timeout=timeout, auth=False) + + # ------------------------------------------------------------------ + # Helpers for subclass RPC methods + + def _get(self, path: str, *, timeout: float = 5.0, auth: bool = True) -> dict: + return request_json( + f"{self.base}{path}", + timeout=timeout, + bearer=self.api_key if auth else None, + ) + + def _post( + self, + path: str, + payload: dict | None = None, + *, + timeout: float = 10.0, + auth: bool = True, + ) -> dict: + return request_json( + f"{self.base}{path}", + method="POST", + payload=payload, + timeout=timeout, + bearer=self.api_key if auth else None, + ) + + def _request( + self, + method: str, + path: str, + payload: dict | None = None, + *, + timeout: float = 10.0, + auth: bool = True, + ) -> dict: + return request_json( + f"{self.base}{path}", + method=method, + payload=payload, + timeout=timeout, + bearer=self.api_key if auth else None, + ) diff --git a/areal/experimental/inference_service/router/__init__.py b/areal/v2/cli/commands/__init__.py similarity index 100% rename from areal/experimental/inference_service/router/__init__.py rename to areal/v2/cli/commands/__init__.py diff --git a/areal/v2/cli/commands/logs.py b/areal/v2/cli/commands/logs.py new file mode 100644 index 0000000000..a4f0ae6703 --- /dev/null +++ b/areal/v2/cli/commands/logs.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""`` logs`` verb factory. + +Every service-style CLI has a ``logs`` command that does the same +thing: resolve the service, find the log directory, tail one component +file. This class packages that pattern. Subcommand CLIs do:: + + from areal.v2.cli.commands.logs import LogsCommand + from areal.v2.cli.inference.lifecycle import inf_lifecycle + + logs_cmd = LogsCommand(lifecycle=inf_lifecycle).build() + inf.add_command(logs_cmd) +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import click + +from areal.v2.cli.lifecycle import ServiceLifecycle + + +class LogsCommand: + default_component: str = "gateway" + default_lines: int = 200 + + def __init__(self, *, lifecycle: ServiceLifecycle) -> None: + self.lifecycle = lifecycle + + def build(self) -> click.Command: + @click.command(name="logs", help="Tail a service log file.") + @click.option("--service", default=None) + @click.option( + "--component", + default=self.default_component, + show_default=True, + help="Component log to tail (gateway / router / worker-N / etc.).", + ) + @click.option("-f", "--follow", is_flag=True, help="Stream appended lines.") + @click.option( + "-n", + "--lines", + type=int, + default=self.default_lines, + show_default=True, + help="Number of trailing lines to show before optional follow.", + ) + def logs_cmd( + service: str | None, component: str, follow: bool, lines: int + ) -> None: + self.execute(service, component, follow, lines) + + return logs_cmd + + def execute( + self, + service: str | None, + component: str, + follow: bool, + lines: int, + ) -> None: + name = self.lifecycle.resolve_service_name(service) + target = self.log_target(name, component) + if not target.exists(): + available = sorted(p.stem for p in target.parent.glob("*.log")) + if not available: + raise click.ClickException(f"no logs found under {target.parent}") + raise click.ClickException( + f"no log named {component!r}; available: {', '.join(available)}" + ) + self._exec_tail(target, lines=lines, follow=follow) + + def log_target(self, service: str, component: str) -> Path: + """Resolve the log file path. Subclasses may override to add a + component name alias / translation layer (e.g. accepting + ``qwen/0/worker`` as shorthand for ``qwen-worker-0``).""" + + return self.lifecycle.store.logs_dir(service) / f"{component}.log" + + def _exec_tail(self, target: Path, *, lines: int, follow: bool) -> None: + cmd = ["tail", f"-n{lines}"] + if follow: + cmd.append("-F") + cmd.append(str(target)) + os.execvp(cmd[0], cmd) diff --git a/areal/v2/cli/config.py b/areal/v2/cli/config.py new file mode 100644 index 0000000000..60abf2b809 --- /dev/null +++ b/areal/v2/cli/config.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""TOML config loader for subcommand CLIs. + +Each CLI's config file lives at ``$AREAL_HOME//config.toml``. +Sections + keys map to ``(verb, click_option_name)`` via a per-CLI +``bindings`` dict that subclasses provide; this base handles all the +flatten + merge + dispatch plumbing. + +Precedence: explicit CLI flag > extra config passed via ``--config`` > +``$AREAL_HOME//config.toml`` > the click option's hard-coded +default. +""" + +from __future__ import annotations + +from pathlib import Path + +from areal.v2.cli.state import NamespacedStateStore + +# (section, key) → (verb_name_or_tuple, click_option_name) +BindingMap = dict[tuple[str, str], tuple[str | tuple[str, ...], str]] + + +class ConfigLoader: + """Subclass and set ``namespace`` + ``bindings`` (or pass them in).""" + + namespace: str = "" + bindings: BindingMap = {} + + def __init__( + self, + *, + namespace: str | None = None, + bindings: BindingMap | None = None, + ) -> None: + if namespace is not None: + self.namespace = namespace + if bindings is not None: + self.bindings = bindings + if not self.namespace: + raise ValueError("ConfigLoader requires a namespace") + self.store = NamespacedStateStore(self.namespace) + + def default_config_path(self) -> Path: + return self.store.config_path() + + def load_click_default_map(self, extra: Path | None = None) -> dict: + """Return a ``click`` ``default_map`` built by merging the + namespace's default config.toml with an optional extra file. + + ``extra`` (typically the ``--config FILE`` flag) wins on + conflicts; bindings absent from the map are silently ignored. + """ + + merged: dict[tuple[str, str], object] = {} + for path in (self.default_config_path(), extra): + if path is None: + continue + merged.update(self._flatten(self._read_toml(path))) + + default_map: dict[str, dict] = {} + for (section, key), value in merged.items(): + binding = self.bindings.get((section, key)) + if binding is None: + continue + verbs, opt = binding + if isinstance(verbs, str): + verbs = (verbs,) + for verb in verbs: + default_map.setdefault(verb, {})[opt] = value + return default_map + + # ------------------------------------------------------------------ + # Helpers — subclasses generally don't need to override these + + def _read_toml(self, path: Path) -> dict: + try: + import tomllib + except ImportError: + return {} + if not path.exists(): + return {} + try: + with open(path, "rb") as f: + return tomllib.load(f) + except Exception: + return {} + + def _flatten(self, toml: dict, parent: str = "") -> dict[tuple[str, str], object]: + """Flatten nested TOML tables to ``(section, key) -> value``. + + Nested keys are dotted on the section side: + ``[register.internal] backend = "..."`` → + ``("register.internal", "backend") -> "..."``. + """ + + out: dict[tuple[str, str], object] = {} + for k, v in toml.items(): + if isinstance(v, dict): + sub_parent = f"{parent}.{k}" if parent else k + for (sec, key), val in self._flatten(v, sub_parent).items(): + out[(sec, key)] = val + else: + out[(parent, k)] = v + return out diff --git a/areal/v2/cli/inference/__init__.py b/areal/v2/cli/inference/__init__.py new file mode 100644 index 0000000000..510e550c93 --- /dev/null +++ b/areal/v2/cli/inference/__init__.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from pathlib import Path + +import click + +from areal.v2.cli.commands.logs import LogsCommand +from areal.v2.cli.inference.commands.deregister import deregister_cmd +from areal.v2.cli.inference.commands.models import models_cmd +from areal.v2.cli.inference.commands.ps import ps_cmd +from areal.v2.cli.inference.commands.register import register_cmd +from areal.v2.cli.inference.commands.run import run_cmd +from areal.v2.cli.inference.commands.status import status_cmd +from areal.v2.cli.inference.commands.stop import stop_cmd +from areal.v2.cli.inference.config import load_click_default_map +from areal.v2.cli.inference.lifecycle import inf_lifecycle + + +@click.group(help="Manage local AReaL inference services.") +@click.option( + "--config", + "config_file", + type=click.Path(exists=True, dir_okay=False, path_type=Path), + default=None, + help="Extra TOML file merged on top of ~/.areal/inf/config.toml.", +) +@click.pass_context +def inf(ctx: click.Context, config_file: Path | None) -> None: + ctx.default_map = load_click_default_map(extra=config_file) + + +inf.add_command(run_cmd) +inf.add_command(stop_cmd) +inf.add_command(status_cmd) +inf.add_command(ps_cmd) +inf.add_command(register_cmd) +inf.add_command(deregister_cmd) +inf.add_command(models_cmd) +inf.add_command(LogsCommand(lifecycle=inf_lifecycle).build()) diff --git a/areal/v2/cli/inference/cli_guide.md b/areal/v2/cli/inference/cli_guide.md new file mode 100644 index 0000000000..5e5e9f078a --- /dev/null +++ b/areal/v2/cli/inference/cli_guide.md @@ -0,0 +1,177 @@ +# AReaL Inference Service CLI + +`areal inf` launches and manages an AReaL inference service on the local machine. It +starts the gateway/router, registers models, inspects service state, and manages logs. + +## Basic concepts + +An inference service typically contains the following components: + +- `gateway`: exposes OpenAI-compatible API and RL API to the outside. +- `router`: maintains the model → worker/data-proxy routing. +- `model worker`: the actual inference backend, e.g. SGLang. +- `data proxy`: records interactions and rewards, and supports trajectory export. + +The CLI's local state is stored under `~/.areal/inf/` by default. The root directory can +be overridden via `AREAL_HOME`: + +```bash +export AREAL_HOME=/path/to/areal-home +``` + +## Launching the service + +Launch an empty inference service: + +```bash +areal inf run \ + --service default \ + --host 127.0.0.1 \ + --port 8080 \ + --admin-api-key areal-admin-key \ + --scheduler local \ + --detach +``` + +`--scheduler` selects the scheduling backend for workers / data-proxies. Only `local` is +supported today (and is the default). Once the service starts, this value is pinned into +the service state, so subsequent `register` / `stop` / `status` calls read it from state +and do not need `--scheduler` again. + +List services known to the local machine: + +```bash +areal inf ps +areal inf status --service default +``` + +`ps` shows the service list; `status` drills into the state of the gateway, router, +data-proxy, workers, etc. + +List registered models: + +```bash +areal inf models --service default +``` + +## Registering a model + +`register` makes the CLI launch a local inference backend together with a data-proxy: + +```bash +areal inf register \ + --service default \ + --model-name qwen-local \ + --backend sglang:d1 \ + --model-path Qwen/Qwen2.5-7B-Instruct \ + --tokenizer-path Qwen/Qwen2.5-7B-Instruct \ + --engine-args "--mem-fraction-static 0.8" \ + --proxy-args "--request-timeout 120 --chat-template-type hf" +``` + +`--engine-args` is a shell-style string forwarded verbatim to the sglang / vllm worker +process; `--proxy-args` is the analogous flag for the data-proxy process. Available +data-proxy flags include `--request-timeout`, `--set-reward-finish-timeout`, +`--tool-call-parser`, `--reasoning-parser`, `--engine-max-tokens`, and +`--chat-template-type {hf|concat}`. + +A model can also be registered directly at `run` time: + +```bash +areal inf run \ + --service default \ + --port 8080 \ + --admin-api-key areal-admin-key \ + --model qwen-local \ + --backend sglang:d1 \ + --model-path Qwen/Qwen2.5-7B-Instruct \ + --engine-args "--mem-fraction-static 0.8" \ + --proxy-args "--request-timeout 120 --chat-template-type hf" \ + --detach +``` + +## Plain inference requests + +Once a model is registered, the gateway's OpenAI-compatible endpoint can be called +directly: + +```bash +curl -sS http://127.0.0.1:8080/v1/chat/completions \ + -H "Authorization: Bearer areal-admin-key" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "qwen-local", + "messages": [ + {"role": "user", "content": "Hi, give me a quick intro to AReaL."} + ], + "max_tokens": 128 + }' +``` + +## Logs and cleanup + +View logs: + +```bash +areal inf logs --service default --component gateway -f +areal inf logs --service default --component router -f +areal inf logs --service default --component qwen-local-worker-0 -f +areal inf logs --service default --component qwen-local-data-proxy-0 -f +``` + +Each model's worker / data-proxy log file is named `-worker-` and +`-data-proxy-`. If `--component` is wrong or the file does not exist, +the CLI prints the available names. + +Deregister a model: + +```bash +areal inf deregister --service default --model-name qwen-local +``` + +Stop the service: + +```bash +areal inf stop --service default +``` + +Force stop: + +```bash +areal inf stop --service default --force +``` + +## Configuration file + +`areal inf` reads a default config from: + +```bash +~/.areal/inf/config.toml +``` + +Additional config files can be passed in: + +```bash +areal inf --config ./inf.toml run --service default --detach +``` + +Example: + +```toml +[default] +service = "default" + +[launch] +gateway_host = "127.0.0.1" +gateway_port = 8080 +routing_strategy = "round_robin" + +[scheduler] +type = "local" + +[register.internal] +backend = "sglang:d1" +model_health_timeout = 600 +engine_args = "--mem-fraction-static 0.8" +proxy_args = "--request-timeout 120 --chat-template-type hf" +``` diff --git a/areal/v2/cli/inference/client.py b/areal/v2/cli/inference/client.py new file mode 100644 index 0000000000..3cb073bdec --- /dev/null +++ b/areal/v2/cli/inference/client.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import Any + +from areal.v2.cli.client import ( + BaseHTTPClient, + ServiceHTTPError, + ServiceUnreachable, +) + +# Legacy aliases — preserve the old names imported by inference subcommands +# (and downstream callers) so this swap is mechanical. New code should use +# the scaffold names directly. +GatewayHTTPError = ServiceHTTPError +GatewayUnreachable = ServiceUnreachable + + +class GatewayClient(BaseHTTPClient): + def list_models(self, *, timeout: float = 5.0) -> dict[str, Any]: + return self._get("/models", timeout=timeout) + + def register_model(self, payload: dict, *, timeout: float = 30.0) -> dict[str, Any]: + return self._post("/register_model", payload=payload, timeout=timeout) + + +class RouterClient(BaseHTTPClient): + def register_worker(self, addr: str, *, timeout: float = 10.0) -> dict[str, Any]: + return self._post("/register", payload={"worker_addr": addr}, timeout=timeout) + + def unregister_worker(self, addr: str, *, timeout: float = 10.0) -> dict[str, Any]: + return self._post("/unregister", payload={"worker_addr": addr}, timeout=timeout) + + def remove_model(self, name: str, *, timeout: float = 10.0) -> dict[str, Any]: + return self._post("/remove_model", payload={"name": name}, timeout=timeout) diff --git a/areal/experimental/inference_service/sglang/__init__.py b/areal/v2/cli/inference/commands/__init__.py similarity index 100% rename from areal/experimental/inference_service/sglang/__init__.py rename to areal/v2/cli/inference/commands/__init__.py diff --git a/areal/v2/cli/inference/commands/deregister.py b/areal/v2/cli/inference/commands/deregister.py new file mode 100644 index 0000000000..332ba7d5f3 --- /dev/null +++ b/areal/v2/cli/inference/commands/deregister.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import click + +from areal.v2.cli.client import ServiceHTTPError, ServiceUnreachable +from areal.v2.cli.inference.client import RouterClient +from areal.v2.cli.inference.common import logger +from areal.v2.cli.inference.lifecycle import inf_lifecycle +from areal.v2.cli.process import kill_pids + + +@click.command(name="deregister", help="Deregister a model and tear down its workers.") +@click.option("--model-name", required=True, help="Model name to deregister.") +@click.option("--service", default=None, help="Target service instance.") +@click.option("--grace", type=float, default=10.0, show_default=True) +@click.option("--force", is_flag=True, help="SIGKILL workers immediately.") +def deregister_cmd( + model_name: str, service: str | None, grace: float, force: bool +) -> None: + raise SystemExit(do_deregister(model_name, grace, force, service=service) or 0) + + +def do_deregister( + model_name: str, grace: float, force: bool, *, service: str | None = None +) -> int: + state = inf_lifecycle.load_running_state(service) + if model_name not in state.models: + raise click.ClickException( + f"model {model_name!r} is not registered in service {state.service!r}" + ) + entry = state.models[model_name] + router = RouterClient(state.router_url, state.admin_api_key) + + try: + router.remove_model(model_name) + except ServiceHTTPError as exc: + if exc.status != 404: + logger.warning("router remove_model %s returned %d", model_name, exc.status) + except ServiceUnreachable as exc: + logger.warning("router unreachable while removing %s: %s", model_name, exc) + + # Router unregister → kill data-proxies → kill workers (same data-flow + # order as terminate_runtime_state). + for r in entry.replicas: + try: + router.unregister_worker(r.data_proxy.addr) + except (ServiceHTTPError, ServiceUnreachable) as exc: + logger.warning("router unregister %s failed: %s", r.data_proxy.addr, exc) + + proxy_pids = [r.data_proxy.pid for r in entry.replicas if r.data_proxy.pid > 0] + worker_pids = [r.worker.pid for r in entry.replicas if r.worker.pid > 0] + effective_grace = 0.0 if force else grace + for pids in (proxy_pids, worker_pids): + if pids: + kill_pids(pids, grace_s=effective_grace) + + del state.model_state.models[model_name] + state.model_state.save() + logger.info("deregistered model %r from service %r", model_name, state.service) + return 0 diff --git a/areal/v2/cli/inference/commands/models.py b/areal/v2/cli/inference/commands/models.py new file mode 100644 index 0000000000..4c44f2b313 --- /dev/null +++ b/areal/v2/cli/inference/commands/models.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from dataclasses import asdict + +import click + +from areal.v2.cli.inference.lifecycle import inf_lifecycle +from areal.v2.cli.utils import json_or_table + + +@click.command(name="models", help="List registered models.") +@click.option("--service", default=None, help="Target service instance.") +@click.option("--json", "as_json", is_flag=True, help="Emit JSON.") +def models_cmd(service: str | None, as_json: bool) -> None: + raise SystemExit(do_models(as_json, service=service) or 0) + + +def do_models(as_json: bool, *, service: str | None = None) -> int: + state = inf_lifecycle.running_state(service) + if state is None: + click.echo("[]" if as_json else "service not running") + return 0 + + payload = [{"name": name, **asdict(entry)} for name, entry in state.models.items()] + json_or_table(payload, as_json=as_json, table_renderer=_print_models_table) + return 0 + + +def _print_models_table(rows: list[dict]) -> None: + if not rows: + click.echo("no models registered") + return + cols = ("NAME", "BACKEND", "WORKERS") + table = [(row["name"], row["backend"], str(len(row["replicas"]))) for row in rows] + widths = [max(len(r[i]) for r in (cols, *table)) for i in range(len(cols))] + fmt = " ".join(f"{{:<{w}}}" for w in widths) + click.echo(fmt.format(*cols)) + for row in table: + click.echo(fmt.format(*row)) diff --git a/areal/v2/cli/inference/commands/ps.py b/areal/v2/cli/inference/commands/ps.py new file mode 100644 index 0000000000..16f003b0b4 --- /dev/null +++ b/areal/v2/cli/inference/commands/ps.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import click + +from areal.v2.cli.inference.lifecycle import inf_lifecycle +from areal.v2.cli.utils import json_or_table + + +@click.command(name="ps", help="List locally known inference services.") +@click.option("--json", "as_json", is_flag=True, help="Emit JSON.") +@click.option("--all", "include_all", is_flag=True, help="Include stale services.") +def ps_cmd(as_json: bool, include_all: bool) -> None: + raise SystemExit(do_ps(as_json, include_all) or 0) + + +def do_ps(as_json: bool, include_all: bool) -> int: + rows: list[dict] = [] + for service in inf_lifecycle.list_services(): + try: + state = inf_lifecycle.load_state(service) + except Exception: + if include_all: + rows.append({"service": service, "status": "stale"}) + continue + running = state.gateway_alive() + if running or include_all: + rows.append( + { + "service": service, + "status": "running" if running else "stale", + "backend": state.backend, + "gateway_url": state.gateway_url, + "gateway_pid": state.gateway_pid, + "router_url": state.router_url, + "models": len(state.models), + } + ) + + json_or_table(rows, as_json=as_json, table_renderer=_print_table) + return 0 + + +def _print_table(rows: list[dict]) -> None: + if not rows: + click.echo("no inference services") + return + + def _pid_cell(row: dict) -> str: + pid = row.get("gateway_pid", 0) + return str(pid) if pid else "-" + + table = [ + ( + row["service"], + row["status"], + row.get("backend", "-"), + str(row.get("models", "")), + row.get("gateway_url", ""), + _pid_cell(row), + ) + for row in rows + ] + cols = ("SERVICE", "STATUS", "BACKEND", "MODELS", "GATEWAY", "PID") + widths = [max(len(str(r[i])) for r in (cols, *table)) for i in range(len(cols))] + fmt = " ".join(f"{{:<{w}}}" for w in widths) + click.echo(fmt.format(*cols)) + for row in table: + click.echo(fmt.format(*row)) diff --git a/areal/v2/cli/inference/commands/register.py b/areal/v2/cli/inference/commands/register.py new file mode 100644 index 0000000000..8961453433 --- /dev/null +++ b/areal/v2/cli/inference/commands/register.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import click + +from areal.v2.cli.inference.client import GatewayClient, RouterClient +from areal.v2.cli.inference.common import ( + ENGINE_ARGS_HELP, + PROXY_ARGS_HELP, + logger, + register_model, +) +from areal.v2.cli.inference.lifecycle import inf_lifecycle +from areal.v2.cli.inference.state import store + + +@click.command(name="register", help="Register a model against a running service.") +@click.option("--model-name", required=True, help="Model name to register.") +@click.option("--service", default=None, help="Target service instance.") +@click.option( + "--backend", + default=None, + help="Backend spec; same grammar as InferenceEngineConfig.backend, " + "e.g. 'sglang:d4', 'vllm:d2t4'.", +) +@click.option("--model-path", default=None) +@click.option("--tokenizer-path", default=None) +@click.option("--engine-args", default="", help=ENGINE_ARGS_HELP) +@click.option("--proxy-args", default="", help=PROXY_ARGS_HELP) +@click.option("--model-health-timeout", type=float, default=600.0, show_default=True) +@click.option( + "--log-level", + type=click.Choice(["debug", "info", "warning", "error"]), + default="info", + show_default=True, +) +def register_cmd(model_name: str, service: str | None, **opts) -> None: + raise SystemExit(do_register(model_name, opts, service=service) or 0) + + +def do_register(model_name: str, opts: dict, *, service: str | None = None) -> int: + service_name = inf_lifecycle.resolve_service_name(service) + with store.lock_model_state(service_name): + state = inf_lifecycle.load_running_state(service_name) + if model_name in state.models: + raise click.ClickException( + f"model {model_name!r} already registered in service {state.service!r}" + ) + + entry = register_model( + model_name=model_name, + opts=opts, + gateway=GatewayClient(state.gateway_url, state.admin_api_key), + router=RouterClient(state.router_url, state.admin_api_key), + log_dir=store.logs_dir(state.service), + admin_api_key=state.admin_api_key, + scheduler_backend=state.backend, + occupied_gpus=state.model_state.occupied_gpus(), + ) + state.model_state.models[model_name] = entry + state.model_state.save() + gpus_used = sorted({g for r in entry.replicas for g in r.worker.gpu_devices}) + logger.info( + "registered model %r in service %r (%d replica(s), GPUs=%s)", + model_name, + state.service, + len(entry.replicas), + gpus_used, + ) + return 0 diff --git a/areal/v2/cli/inference/commands/run.py b/areal/v2/cli/inference/commands/run.py new file mode 100644 index 0000000000..5342abd2ad --- /dev/null +++ b/areal/v2/cli/inference/commands/run.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import time + +import click + +from areal.v2.cli.inference.client import GatewayClient, RouterClient +from areal.v2.cli.inference.common import ( + ENGINE_ARGS_HELP, + PROXY_ARGS_HELP, + logger, + register_model, + terminate_runtime_state, + validate_register_opts, +) +from areal.v2.cli.inference.launcher import spawn_gateway, spawn_router +from areal.v2.cli.inference.lifecycle import inf_lifecycle +from areal.v2.cli.inference.scheduler import TaskHandle +from areal.v2.cli.inference.state import ( + DEFAULT_SERVICE, + ModelState, + RuntimeState, + ServiceState, + store, +) +from areal.v2.cli.process import kill_pids +from areal.v2.cli.utils import wait_client_health +from areal.v2.cli.watcher import ForegroundWatcher + + +@click.command(name="run", help="Start an inference service (gateway + router).") +@click.option("--service", default=DEFAULT_SERVICE, show_default=True) +@click.option("--port", type=int, default=8080, show_default=True, help="Gateway port.") +@click.option( + "--host", default="127.0.0.1", show_default=True, help="Gateway bind host." +) +@click.option("--admin-api-key", default="areal-admin-key", show_default=True) +@click.option( + "--routing-strategy", + type=click.Choice(["round_robin", "least_busy"]), + default="round_robin", + show_default=True, +) +@click.option( + "--log-level", + type=click.Choice(["debug", "info", "warning", "error"]), + default="info", + show_default=True, +) +@click.option( + "--launch-timeout", + type=float, + default=30.0, + show_default=True, + help="Seconds to wait for gateway /health.", +) +@click.option("-d", "--detach", is_flag=True, help="Fork the daemon and exit.") +@click.option("--model", default=None, help="Register this model at startup.") +@click.option( + "--backend", + default=None, + help="Backend spec; same grammar as InferenceEngineConfig.backend, " + "e.g. 'sglang:d4', 'vllm:d2t4'.", +) +@click.option("--model-path", default=None) +@click.option("--tokenizer-path", default=None) +@click.option("--engine-args", default="", show_default=False, help=ENGINE_ARGS_HELP) +@click.option("--proxy-args", default="", show_default=False, help=PROXY_ARGS_HELP) +@click.option( + "--model-health-timeout", + type=float, + default=600.0, + show_default=True, + help="Seconds to wait for the model server to come up.", +) +@click.option( + "--scheduler", + type=click.Choice(["local"]), + default=None, + help=( + "Scheduler backend used to place workers and data-proxies. Pinned " + "for the service's lifetime — register / stop / status will read it " + "from saved state. Defaults to [scheduler].type in config.toml, " + "falling back to 'local'." + ), +) +@click.option("--force", is_flag=True, help="Replace stale or running service state.") +def run_cmd(**opts) -> None: + raise SystemExit(do_run(opts) or 0) + + +def do_run(opts: dict) -> int: + service = opts["service"] + if opts["force"]: + inf_lifecycle.force_replace_slot(service, grace_s=5.0) + else: + inf_lifecycle.refuse_if_running(service) + + if opts["model"]: + validate_register_opts(opts) + elif opts["backend"]: + raise click.UsageError("model registration flags require --model.") + + log_dir = store.logs_dir(service) + logger.info("starting inference service %r (logs: %s)", service, log_dir) + + router_pid, router_port = spawn_router( + host="127.0.0.1", + admin_api_key=opts["admin_api_key"], + routing_strategy=opts["routing_strategy"], + log_level=opts["log_level"], + log_file=log_dir / "router.log", + ) + router_url = f"http://127.0.0.1:{router_port}" + logger.info("router pid=%d %s", router_pid, router_url) + + time.sleep(0.3) + gateway_pid = spawn_gateway( + host=opts["host"], + port=opts["port"], + admin_api_key=opts["admin_api_key"], + router_url=router_url, + log_level=opts["log_level"], + log_file=log_dir / "gateway.log", + ) + host_for_url = "127.0.0.1" if opts["host"] in ("0.0.0.0", "::") else opts["host"] + gateway_url = f"http://{host_for_url}:{opts['port']}" + logger.info("gateway pid=%d %s", gateway_pid, gateway_url) + + gateway_client = GatewayClient(gateway_url, opts["admin_api_key"]) + router_client = RouterClient(router_url, opts["admin_api_key"]) + try: + wait_client_health( + gateway_client, timeout=opts["launch_timeout"], label="gateway" + ) + except BaseException: + kill_pids([gateway_pid, router_pid], grace_s=5.0) + raise + + backend = opts.get("scheduler") or "local" + + service_state = ServiceState( + service=service, + backend=backend, + gateway_handle=TaskHandle( + host=host_for_url, + ports=[opts["port"]], + gpu_devices=[], + ref={"pid": gateway_pid}, + ), + router_handle=TaskHandle( + host="127.0.0.1", + ports=[router_port], + gpu_devices=[], + ref={"pid": router_pid}, + ), + admin_api_key=opts["admin_api_key"], + started_at=time.time(), + ) + model_state = ModelState(service=service) + with store.lock_model_state(service): + service_state.save() + model_state.save() + + if opts["model"]: + try: + entry = register_model( + model_name=opts["model"], + opts=opts, + gateway=gateway_client, + router=router_client, + log_dir=log_dir, + admin_api_key=opts["admin_api_key"], + scheduler_backend=backend, + occupied_gpus=model_state.occupied_gpus(), + ) + model_state.models[opts["model"]] = entry + model_state.save() + except BaseException: + _cleanup_runtime(service, service_state, model_state, grace_s=5.0) + raise + + logger.info("service %r ready pid=%d url=%s", service, gateway_pid, gateway_url) + if opts["model"]: + logger.info("default model: %s (%s)", opts["model"], opts["backend"]) + + if opts["detach"]: + return 0 + + logger.info( + "foreground watcher running. Ctrl+C / SIGTERM tears down the service; " + "closing the terminal (SIGHUP) detaches the watcher and leaves it " + "running — use `areal inf stop --service %s` to terminate.", + service, + ) + watcher = ForegroundWatcher( + is_alive=lambda: inf_lifecycle.gateway_alive( + RuntimeState(service_state=service_state, model_state=model_state) + ), + teardown=lambda: _cleanup_runtime( + service, service_state, model_state, grace_s=10.0 + ), + service_name=service, + ) + rc = watcher.watch() + # gateway died externally (is_alive flipped to False on its own) — the + # children are gone, just clean the state files. + if rc == 0 and not inf_lifecycle.gateway_alive( + RuntimeState(service_state=service_state, model_state=model_state) + ): + ServiceState.remove(service) + ModelState.remove(service) + return rc + + +def _cleanup_runtime( + service: str, + service_state: ServiceState, + model_state: ModelState, + *, + grace_s: float, +) -> None: + model_state = _latest_model_state_for_cleanup(service, model_state) + try: + terminate_runtime_state( + RuntimeState(service_state=service_state, model_state=model_state), + grace_s=grace_s, + ) + finally: + ServiceState.remove(service) + ModelState.remove(service) + + +def _latest_model_state_for_cleanup( + service: str, + fallback: ModelState, +) -> ModelState: + try: + latest = ModelState.load(service) + except Exception: + return fallback + latest.models.update(fallback.models) + return latest diff --git a/areal/v2/cli/inference/commands/status.py b/areal/v2/cli/inference/commands/status.py new file mode 100644 index 0000000000..ea9e691bd4 --- /dev/null +++ b/areal/v2/cli/inference/commands/status.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from dataclasses import asdict + +import click + +from areal.v2.cli.inference.common import ( + format_gpu_count, + format_placement, + format_ref, +) +from areal.v2.cli.inference.lifecycle import inf_lifecycle +from areal.v2.cli.inference.scheduler import TaskHandle +from areal.v2.cli.inference.state import RuntimeState +from areal.v2.cli.state import SupportsComponentProbe +from areal.v2.cli.status import ColumnSpec, StatusReporter + + +@click.command( + name="status", + help="Show inference service status — per-component drill-down for one service.", +) +@click.option("--service", default=None, help="Target service instance.") +@click.option("--json", "as_json", is_flag=True, help="Emit JSON.") +def status_cmd(service: str | None, as_json: bool) -> None: + raise SystemExit(do_status(as_json, service=service) or 0) + + +def do_status(as_json: bool, *, service: str | None = None) -> int: + service_name = inf_lifecycle.resolve_service_name(service) + if not inf_lifecycle.state_path(service_name).exists(): + if as_json: + click.echo( + json.dumps({"service": service_name, "running": False}, indent=2) + ) + else: + click.echo(f"service {service_name!r} not running") + return 0 + + try: + state = inf_lifecycle.load_state(service_name) + except Exception as exc: + raise click.ClickException(f"failed to load state: {exc}") from exc + + components = list(state.components()) + reporter = StatusReporter(components, _columns(state.backend)) + alive = reporter.probe_all() + + if as_json: + click.echo(json.dumps(_json_snapshot(state, reporter, alive), indent=2)) + return 0 + + gateway_addr = state.gateway_handle.addr or "-" + reporter.print_table( + reporter.render_rows(alive), + header_line=( + f"service: {state.service} " + f"backend: {state.backend} " + f"gateway: {gateway_addr} " + f"models: {len(state.models)} " + f"running: {'yes' if state.gateway_alive() else 'no'}" + ), + ) + return 0 + + +def _columns(backend: str) -> list[ColumnSpec]: + def _component(label: str, _: SupportsComponentProbe, __: bool) -> str: + return label + + def _placement(_: str, handle: SupportsComponentProbe, __: bool) -> str: + # SupportsComponentProbe doesn't promise the TaskHandle surface, + # but in practice every inf component handle IS a TaskHandle. + return format_placement(backend, handle) or "-" # type: ignore[arg-type] + + def _gpus(_: str, handle: SupportsComponentProbe, __: bool) -> str: + return format_gpu_count(handle) # type: ignore[arg-type] + + def _addr(_: str, handle: SupportsComponentProbe, __: bool) -> str: + return handle.addr or "-" + + def _ref(_: str, handle: SupportsComponentProbe, __: bool) -> str: + return format_ref(backend, handle) # type: ignore[arg-type] + + def _alive(_: str, __: SupportsComponentProbe, alive: bool) -> str: + return "yes" if alive else "no" + + return [ + ColumnSpec("COMPONENT", _component), + ColumnSpec("PLACEMENT", _placement), + ColumnSpec("GPUS", _gpus), + ColumnSpec("ADDR", _addr), + ColumnSpec("REF", _ref), + ColumnSpec("ALIVE", _alive), + ] + + +def _json_snapshot( + state: RuntimeState, reporter: StatusReporter, alive: list[bool] +) -> dict: + handle_pairs: list[tuple[str, TaskHandle]] = list(state.components()) # type: ignore[arg-type] + rows = [ + {**snap, "label": label} + for snap, (label, _) in zip( + reporter.json_snapshot(alive), handle_pairs, strict=True + ) + ] + return { + "service": state.service, + "backend": state.backend, + "running": state.gateway_alive(), + "gateway_handle": asdict(state.gateway_handle), + "router_handle": asdict(state.router_handle), + "started_at": state.started_at, + "models": {name: asdict(entry) for name, entry in state.models.items()}, + "components": rows, + } diff --git a/areal/v2/cli/inference/commands/stop.py b/areal/v2/cli/inference/commands/stop.py new file mode 100644 index 0000000000..88d0d78244 --- /dev/null +++ b/areal/v2/cli/inference/commands/stop.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import click + +from areal.v2.cli.inference.common import logger, terminate_runtime_state +from areal.v2.cli.inference.lifecycle import inf_lifecycle +from areal.v2.cli.inference.state import ModelState, ServiceState + + +@click.command(name="stop", help="Stop an inference service.") +@click.option("--service", default=None, help="Target service instance.") +@click.option("--grace", type=float, default=10.0, show_default=True) +@click.option("--force", is_flag=True, help="SIGKILL immediately.") +@click.option("--keep-state", is_flag=True, help="Keep service state after stopping.") +def stop_cmd(service: str | None, grace: float, force: bool, keep_state: bool) -> None: + raise SystemExit(do_stop(grace, force, service=service, keep_state=keep_state) or 0) + + +def do_stop( + grace: float, + force: bool, + *, + service: str | None = None, + keep_state: bool = False, +) -> int: + service_name = inf_lifecycle.resolve_service_name(service) + if not inf_lifecycle.state_path(service_name).exists(): + click.echo(f"service {service_name!r} not running") + return 0 + + try: + state = inf_lifecycle.load_state(service_name) + except Exception: + logger.warning("stale state; removing") + if not keep_state: + ServiceState.remove(service_name) + ModelState.remove(service_name) + return 0 + + terminate_runtime_state(state, grace_s=grace, force=force) + + if not keep_state: + ServiceState.remove(service_name) + ModelState.remove(service_name) + click.echo(f"service {service_name!r} stopped") + return 0 diff --git a/areal/v2/cli/inference/common.py b/areal/v2/cli/inference/common.py new file mode 100644 index 0000000000..72b9c3d203 --- /dev/null +++ b/areal/v2/cli/inference/common.py @@ -0,0 +1,335 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Domain-specific helpers for the inference CLI. + +Lifecycle predicates, HTTP polling, output dispatch, and logger setup +live in the shared scaffold (`areal.v2.cli.lifecycle`, +`.utils`, `.client`, `.status`). This module keeps only the +inference-specific pieces: backend spec parsing, model registration, +and TaskHandle column formatters. +""" + +from __future__ import annotations + +from pathlib import Path + +import click + +from areal.api.alloc_mode import ModelAllocation +from areal.v2.cli.client import ServiceHTTPError, ServiceUnreachable +from areal.v2.cli.inference.client import GatewayClient, RouterClient +from areal.v2.cli.inference.launcher import ( + build_data_proxy_task_spec, + build_sglang_task_spec, + build_vllm_task_spec, +) +from areal.v2.cli.inference.scheduler import ( + Scheduler, + TaskHandle, + build_scheduler, +) +from areal.v2.cli.inference.state import ( + ModelEntry, + ModelReplica, + RuntimeState, +) +from areal.v2.cli.process import kill_pids +from areal.v2.cli.utils import register_cli_logger, wait_http_health + +logger = register_cli_logger("InfCli") + + +ENGINE_ARGS_HELP = ( + "Shell-style string forwarded verbatim to the sglang / vllm process. " + "Common sglang knobs: --mem-fraction-static 0.85, " + "--max-running-requests 256, --chunked-prefill-size 4096, " + "--disable-radix-cache, --enable-torch-compile. " + "See the sglang / vllm CLI docs for the full surface." +) + +PROXY_ARGS_HELP = ( + "Shell-style string forwarded verbatim to the data-proxy process. " + "Available flags: --request-timeout, --set-reward-finish-timeout, " + "--tool-call-parser, --reasoning-parser, --engine-max-tokens, " + "--chat-template-type {hf|concat}." +) + + +def parse_backend_spec(spec: str) -> tuple[str, int, int, int]: + """Parse a backend spec into ``(engine, tp, dp, pp)``. + + Delegates to ``ModelAllocation.from_str`` so the CLI accepts the same + grammar as ``InferenceEngineConfig.backend`` in YAML configs — + ``"sglang:d4"``, ``"vllm:d2t4"``, etc. The CLI restricts the engine + set to the two inference backends it can spawn locally. + """ + + try: + alloc = ModelAllocation.from_str(spec) + except Exception as exc: + raise click.BadParameter(f"invalid --backend spec {spec!r}: {exc}") from exc + if alloc.backend not in ("sglang", "vllm"): + raise click.BadParameter( + f"backend must be one of: sglang, vllm; got {alloc.backend!r}" + ) + return ( + alloc.backend, + alloc.parallel.tensor_parallel_size, + alloc.parallel.data_parallel_size, + alloc.parallel.pipeline_parallel_size, + ) + + +def split_args(value: str) -> list[str]: + import shlex + + return shlex.split(value) if value else [] + + +def terminate_runtime_state( + state: RuntimeState, *, grace_s: float, force: bool = False +) -> None: + """Tear down the whole inference service in data-flow order: kill the + data-proxies first (so in-flight requests fail at the proxy boundary + rather than mid-worker), then workers, then gateway, then router. + + ``force=True`` skips the grace period entirely — equivalent to + ``kill_pids(pids, grace_s=0)`` for every phase. + """ + + def _pids(handles: list[TaskHandle]) -> list[int]: + return [h.pid for h in handles if h.pid > 0] + + phases = ( + _pids(state.model_state.all_data_proxies()), + _pids(state.model_state.all_workers()), + _pids([state.gateway_handle]), + _pids([state.router_handle]), + ) + effective_grace = 0.0 if force else grace_s + for pids in phases: + if not pids: + continue + kill_pids(pids, grace_s=effective_grace) + + +def format_placement(backend: str, handle: TaskHandle) -> str: + if backend == "local": + import socket + + return socket.gethostname() + if backend == "k8s": + node = handle.ref.get("node", "?") + pod = handle.ref.get("pod_name", "?") + return f"{node}/{pod}" + if backend == "slurm": + node = handle.ref.get("node", "?") + job = handle.ref.get("job_id", "?") + return f"{node}/job={job}" + return "-" + + +def format_ref(backend: str, handle: TaskHandle) -> str: + if backend == "local": + return f"pid={handle.pid}" if handle.pid > 0 else "-" + if backend == "k8s": + pod = handle.ref.get("pod_name", "") + return f"pod={pod}" if pod else "-" + if backend == "slurm": + job = handle.ref.get("job_id", "") + return f"job={job}" if job else "-" + return "-" + + +def format_gpu_count(handle: TaskHandle) -> str: + n = len(handle.gpu_devices) + return f"×{n}" if n > 0 else "-" + + +def register_internal( + *, + model: str, + backend: str, + model_path: str, + tokenizer_path: str, + engine_extra: list[str], + proxy_extra: list[str], + model_health_timeout: float, + log_level: str, + admin_api_key: str, + gateway: GatewayClient, + router: RouterClient, + log_dir: Path, + scheduler: Scheduler, +) -> list[ModelReplica]: + engine, tp, dp, pp = parse_backend_spec(backend) + if pp > 1: + raise click.ClickException( + "pp > 1 is not supported by `areal inf` (single-node only)." + ) + + replicas: list[ModelReplica] = [] + spawned_handles: list[TaskHandle] = [] + + try: + for rank in range(dp): + worker_log = log_dir / f"{model}-worker-{rank}.log" + if engine == "sglang": + worker_spec = build_sglang_task_spec( + name=f"worker/{model}/{rank}", + model_path=model_path, + tp=tp, + extra_args=engine_extra, + log_file=worker_log, + ) + else: + worker_spec = build_vllm_task_spec( + name=f"worker/{model}/{rank}", + model_path=model_path, + tp=tp, + pp=pp, + extra_args=engine_extra, + log_file=worker_log, + ) + worker_handle = scheduler.submit(worker_spec) + spawned_handles.append(worker_handle) + logger.info( + "spawned %s worker %d/%d pid=%d port=%d gpus=%s", + engine, + rank, + dp, + worker_handle.pid, + worker_handle.ports[0], + worker_handle.gpu_devices, + ) + wait_http_health( + worker_handle.addr, + pid=worker_handle.pid, + timeout=model_health_timeout, + label=f"{engine} worker {rank}", + poll_interval=1.0, + ) + + proxy_log = log_dir / f"{model}-data-proxy-{rank}.log" + proxy_spec = build_data_proxy_task_spec( + name=f"data_proxy/{model}/{rank}", + backend_addr=worker_handle.addr, + backend_type=engine, + tokenizer_path=tokenizer_path, + admin_api_key=admin_api_key, + log_level=log_level, + extra_args=proxy_extra, + log_file=proxy_log, + ) + proxy_handle = scheduler.submit(proxy_spec) + spawned_handles.append(proxy_handle) + logger.info( + "spawned data-proxy %d/%d pid=%d port=%d", + rank, + dp, + proxy_handle.pid, + proxy_handle.ports[0], + ) + wait_http_health( + proxy_handle.addr, + pid=proxy_handle.pid, + timeout=30.0, + label=f"data-proxy {rank}", + poll_interval=1.0, + ) + + replicas.append(ModelReplica(data_proxy=proxy_handle, worker=worker_handle)) + + proxy_addrs = [r.data_proxy.addr for r in replicas] + for addr in proxy_addrs: + try: + router.register_worker(addr) + except (ServiceUnreachable, ServiceHTTPError) as exc: + raise click.ClickException( + f"router register_worker {addr} failed: {exc}" + ) from exc + + try: + gateway.register_model( + { + "model": model, + "url": "", + "api_key": "", + "data_proxy_addrs": proxy_addrs, + } + ) + except (ServiceUnreachable, ServiceHTTPError) as exc: + raise click.ClickException(f"gateway register_model failed: {exc}") from exc + + except BaseException: + if spawned_handles: + logger.error( + "internal register failed; killing %d spawned worker(s)", + len(spawned_handles), + ) + pids = [h.pid for h in spawned_handles if h.pid > 0] + kill_pids(pids, grace_s=10.0) + raise + + return replicas + + +def validate_register_opts(opts: dict) -> None: + if not opts.get("backend"): + raise click.UsageError("--backend is required.") + if not opts.get("model_path"): + raise click.UsageError("--backend requires --model-path .") + + +def register_model( + *, + model_name: str, + opts: dict, + gateway: GatewayClient, + router: RouterClient, + log_dir: Path, + admin_api_key: str, + scheduler_backend: str, + occupied_gpus: set[int], +) -> ModelEntry: + """Validate registration opts and spawn the model's worker + data-proxy + fleet. The returned ModelEntry is meant to be slotted into model_state + by the caller — this helper does not touch model_state itself so the + caller can keep the file lock scope explicit.""" + + validate_register_opts(opts) + scheduler = build_scheduler(scheduler_backend, occupied_gpus=occupied_gpus) + replicas = register_internal( + model=model_name, + backend=opts["backend"], + model_path=opts["model_path"], + tokenizer_path=opts.get("tokenizer_path") or opts["model_path"], + engine_extra=split_args(opts.get("engine_args", "")), + proxy_extra=split_args(opts.get("proxy_args", "")), + model_health_timeout=opts.get("model_health_timeout", 600.0), + log_level=opts.get("log_level", "info"), + admin_api_key=admin_api_key, + gateway=gateway, + router=router, + log_dir=log_dir, + scheduler=scheduler, + ) + return ModelEntry(backend=opts["backend"], replicas=replicas) + + +# Imports kept at bottom: these are stdlib re-exports preserved for +# downstream callers that still import them via this module. +__all__ = [ + "ENGINE_ARGS_HELP", + "PROXY_ARGS_HELP", + "format_gpu_count", + "format_placement", + "format_ref", + "logger", + "parse_backend_spec", + "register_internal", + "register_model", + "split_args", + "terminate_runtime_state", + "validate_register_opts", +] diff --git a/areal/v2/cli/inference/config.py b/areal/v2/cli/inference/config.py new file mode 100644 index 0000000000..d28d325bb9 --- /dev/null +++ b/areal/v2/cli/inference/config.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""User config loader for ``areal inf``. + +~/.areal/inf/config.toml overrides built-in CLI defaults. Sections map +to verbs: + + [default] admin_api_key, log_level (applied to all verbs) + [launch] run-time gateway / router / strategy / timeouts + [scheduler] scheduler backend selection (applied to run only) + [register.internal] register / inline-register defaults + +Precedence (highest first): + 1. explicit CLI flag + 2. config.toml + 3. hard-coded default in the click @option + +Missing / malformed file is treated as empty — never crashes the CLI. +""" + +from __future__ import annotations + +from pathlib import Path + +from areal.v2.cli.config import BindingMap, ConfigLoader +from areal.v2.cli.inference.state import INF_NAMESPACE + +INF_BINDINGS: BindingMap = { + # [default] + ( + "default", + "service", + ): ( + ( + "run", + "stop", + "status", + "register", + "deregister", + "models", + "logs", + ), + "service", + ), + ("default", "admin_api_key"): ("run", "admin_api_key"), + ("default", "log_level"): ("run", "log_level"), + # [launch] — applied to `run` + ("launch", "gateway_host"): ("run", "host"), + ("launch", "gateway_port"): ("run", "port"), + ("launch", "routing_strategy"): ("run", "routing_strategy"), + ("launch", "launch_timeout"): ("run", "launch_timeout"), + # [scheduler] — applied to `run` only; other verbs read from ServiceState + ("scheduler", "type"): ("run", "scheduler"), + # [register.internal] — applied to `register` + ("register.internal", "backend"): ("register", "backend"), + ("register.internal", "model_health_timeout"): ("register", "model_health_timeout"), + ("register.internal", "engine_args"): ("register", "engine_args"), + ("register.internal", "proxy_args"): ("register", "proxy_args"), +} + + +inf_config_loader = ConfigLoader(namespace=INF_NAMESPACE, bindings=INF_BINDINGS) + + +def load_click_default_map(extra: Path | None = None) -> dict: + return inf_config_loader.load_click_default_map(extra=extra) diff --git a/areal/v2/cli/inference/launcher.py b/areal/v2/cli/inference/launcher.py new file mode 100644 index 0000000000..19a104a4db --- /dev/null +++ b/areal/v2/cli/inference/launcher.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import sys +from pathlib import Path + +from areal.utils.network import find_free_ports +from areal.v2.cli.inference.scheduler import ( + TaskAllocation, + TaskResources, + TaskSpec, +) +from areal.v2.cli.process import spawn_process + +# Gateway and router are service-level singletons spawned directly via +# spawn_process — not submitted to a Scheduler — so their lifecycle is +# anchored to the CLI's state files regardless of scheduler backend. + + +def spawn_router( + *, + host: str, + admin_api_key: str, + routing_strategy: str, + log_level: str, + log_file: Path, +) -> tuple[int, int]: + port = find_free_ports(1)[0] + cmd = [ + sys.executable, + "-m", + "areal.v2.inference_service.router", + "--host", + host, + "--port", + str(port), + "--admin-api-key", + admin_api_key, + "--routing-strategy", + routing_strategy, + "--log-level", + log_level, + ] + return spawn_process(cmd, log_file), port + + +def spawn_gateway( + *, + host: str, + port: int, + admin_api_key: str, + router_url: str, + log_level: str, + log_file: Path, +) -> int: + cmd = [ + sys.executable, + "-m", + "areal.v2.inference_service.gateway", + "--host", + host, + "--port", + str(port), + "--admin-api-key", + admin_api_key, + "--router-addr", + router_url, + "--log-level", + log_level, + ] + return spawn_process(cmd, log_file) + + +# sglang / vllm builders always pass base_gpu_id=0; the scheduler masks +# physical GPUs via CUDA_VISIBLE_DEVICES so the launched process sees +# its devices as 0..N-1. + + +def build_sglang_task_spec( + *, + name: str, + model_path: str, + tp: int, + extra_args: list[str], + log_file: Path, +) -> TaskSpec: + def cmd_builder(alloc: TaskAllocation) -> list[str]: + from areal.api.cli_args import SGLangConfig + + cfg = SGLangConfig(model_path=model_path, log_requests=True) + cmd = list( + SGLangConfig.build_cmd( + sglang_config=cfg, + tp_size=tp, + base_gpu_id=0, + host=alloc.host, + port=alloc.ports[0], + n_nodes=1, + node_rank=0, + pp_size=1, + ) + ) + cmd.extend(extra_args) + return cmd + + return TaskSpec( + name=name, + cmd_builder=cmd_builder, + log_file=log_file, + resources=TaskResources(gpu=tp, ports=1), + ) + + +def build_vllm_task_spec( + *, + name: str, + model_path: str, + tp: int, + pp: int, + extra_args: list[str], + log_file: Path, +) -> TaskSpec: + def cmd_builder(alloc: TaskAllocation) -> list[str]: + from areal.api.cli_args import vLLMConfig + + cfg = vLLMConfig(model=model_path) + cmd = list( + vLLMConfig.build_cmd( + vllm_config=cfg, + tp_size=tp, + pp_size=pp, + host=alloc.host, + port=alloc.ports[0], + ) + ) + cmd.extend(extra_args) + return cmd + + return TaskSpec( + name=name, + cmd_builder=cmd_builder, + log_file=log_file, + resources=TaskResources(gpu=tp * pp, ports=1), + ) + + +def build_data_proxy_task_spec( + *, + name: str, + backend_addr: str, + backend_type: str, + tokenizer_path: str, + admin_api_key: str, + log_level: str, + extra_args: list[str], + log_file: Path, +) -> TaskSpec: + def cmd_builder(alloc: TaskAllocation) -> list[str]: + cmd = [ + sys.executable, + "-m", + "areal.v2.inference_service.data_proxy", + "--host", + alloc.host, + "--port", + str(alloc.ports[0]), + "--backend-addr", + backend_addr, + "--backend-type", + backend_type, + "--tokenizer-path", + tokenizer_path, + "--admin-api-key", + admin_api_key, + "--log-level", + log_level, + ] + cmd.extend(extra_args) + return cmd + + return TaskSpec( + name=name, + cmd_builder=cmd_builder, + log_file=log_file, + resources=TaskResources(gpu=0, ports=1), + ) diff --git a/areal/v2/cli/inference/lifecycle.py b/areal/v2/cli/inference/lifecycle.py new file mode 100644 index 0000000000..4856600374 --- /dev/null +++ b/areal/v2/cli/inference/lifecycle.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from areal.v2.cli.inference.state import ( + INF_NAMESPACE, + ModelState, + RuntimeState, + ServiceState, + store, +) +from areal.v2.cli.lifecycle import ServiceLifecycle +from areal.v2.cli.process import kill_pids + + +class InferenceLifecycle(ServiceLifecycle): + """Adds the inf-specific behavior on top of scaffold's lifecycle: + + - State spans two files (service + model registry). ``force_replace_slot`` + walks the inf raw-JSON helper which knows about both files, and + removes both on cleanup. + """ + + def force_replace_slot(self, service: str, *, grace_s: float = 5.0) -> None: + path = self.state_path(service) + if not path.exists(): + return + pids: list[int] = [] + try: + state = self.load_state(service) + pids = self._collect_pids(state) + except Exception: + try: + pids = store.recover_pids_from_raw_state(service) + except Exception: + pids = [] + if pids: + kill_pids(pids, grace_s=grace_s) + ServiceState.remove(service) + ModelState.remove(service) + + +inf_lifecycle = InferenceLifecycle( + namespace=INF_NAMESPACE, + state_class=RuntimeState, + stop_command="areal inf stop", +) diff --git a/areal/v2/cli/inference/scheduler/__init__.py b/areal/v2/cli/inference/scheduler/__init__.py new file mode 100644 index 0000000000..49c8d5cbb5 --- /dev/null +++ b/areal/v2/cli/inference/scheduler/__init__.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 + +from areal.v2.cli.inference.scheduler.base import ( + Scheduler, + SchedulerError, + TaskAllocation, + TaskHandle, + TaskResources, + TaskSpec, + build_scheduler, +) + +__all__ = [ + "Scheduler", + "SchedulerError", + "TaskAllocation", + "TaskHandle", + "TaskResources", + "TaskSpec", + "build_scheduler", +] diff --git a/areal/v2/cli/inference/scheduler/base.py b/areal/v2/cli/inference/scheduler/base.py new file mode 100644 index 0000000000..727fd37d9c --- /dev/null +++ b/areal/v2/cli/inference/scheduler/base.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Callable +from dataclasses import dataclass, field +from pathlib import Path + + +class SchedulerError(Exception): + pass + + +@dataclass +class TaskResources: + gpu: int = 0 + ports: int = 1 + + +@dataclass +class TaskAllocation: + host: str + ports: list[int] + gpu_devices: list[int] + + +@dataclass +class TaskSpec: + """cmd_builder receives the TaskAllocation so host/port/GPU can be filled + in at placement time — required for K8s/Slurm where the host is only + known after scheduling.""" + + name: str + cmd_builder: Callable[[TaskAllocation], list[str]] + log_file: Path + resources: TaskResources = field(default_factory=TaskResources) + env: dict[str, str] = field(default_factory=dict) + + +@dataclass +class TaskHandle: + """Backend-agnostic identity of a running task. + + ``ref`` carries backend-specific identity payload — ``{"pid": int}`` for + Local, ``{"pod_name", "namespace"}`` for K8s, ``{"job_id"}`` for Slurm. + """ + + host: str + ports: list[int] + gpu_devices: list[int] = field(default_factory=list) + ref: dict = field(default_factory=dict) + + @property + def addr(self) -> str: + if not self.host or not self.ports: + return "" + return f"http://{self.host}:{self.ports[0]}" + + @property + def pid(self) -> int: + return int(self.ref.get("pid", 0) or 0) + + +class Scheduler(ABC): + @abstractmethod + def submit(self, spec: TaskSpec) -> TaskHandle: ... + + +def build_scheduler(backend: str, **kwargs) -> Scheduler: + if backend == "local": + from areal.v2.cli.inference.scheduler.local import LocalScheduler + + return LocalScheduler(**kwargs) + raise SchedulerError(f"unknown scheduler backend: {backend!r}") diff --git a/areal/v2/cli/inference/scheduler/local.py b/areal/v2/cli/inference/scheduler/local.py new file mode 100644 index 0000000000..3a2c2dfc60 --- /dev/null +++ b/areal/v2/cli/inference/scheduler/local.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os + +from areal.utils.logging import getLogger +from areal.utils.network import find_free_ports +from areal.v2.cli.inference.scheduler.base import ( + Scheduler, + SchedulerError, + TaskAllocation, + TaskHandle, + TaskSpec, +) +from areal.v2.cli.process import spawn_process + +logger = getLogger("InfLocalScheduler") + + +def _detect_gpus() -> list[int]: + vis = os.environ.get("CUDA_VISIBLE_DEVICES") + if vis: + try: + return [int(x) for x in vis.split(",") if x.strip()] + except ValueError: + logger.warning( + "ignoring malformed CUDA_VISIBLE_DEVICES=%r, falling back to detection", + vis, + ) + try: + nvidia = [ + d + for d in os.listdir("/dev") + if d.startswith("nvidia") and d[len("nvidia") :].isdigit() + ] + except OSError: + return [0] + return sorted(int(d[len("nvidia") :]) for d in nvidia) or [0] + + +class LocalScheduler(Scheduler): + """Ephemeral local subprocess scheduler for ``areal inf``. + + The instance is discarded after each CLI command — Popen handles are + not retained, only the PID written into the returned TaskHandle. GPUs + are masked via ``CUDA_VISIBLE_DEVICES`` so the spawned process always + sees its devices as ``0..N-1``, matching K8s device-plugin semantics. + """ + + def __init__( + self, + *, + all_gpus: list[int] | None = None, + occupied_gpus: set[int] | None = None, + host: str = "127.0.0.1", + ) -> None: + self.all_gpus = list(all_gpus) if all_gpus is not None else _detect_gpus() + self.occupied_gpus: set[int] = set(occupied_gpus or ()) + self.host = host + # Ports already handed out by this scheduler instance — passed as + # exclude_ports so the next find_free_ports call cannot return the + # same port between subsequent submits within one CLI command. + self._allocated_ports: set[int] = set() + + def submit(self, spec: TaskSpec) -> TaskHandle: + gpus = self._allocate_gpus(spec.resources.gpu) + ports = self._allocate_ports(spec.resources.ports) + alloc = TaskAllocation(host=self.host, ports=ports, gpu_devices=gpus) + cmd = spec.cmd_builder(alloc) + + env = dict(spec.env) + if gpus: + env["CUDA_VISIBLE_DEVICES"] = ",".join(str(g) for g in gpus) + + try: + pid = spawn_process(cmd, spec.log_file, env=env) + except OSError as exc: + raise SchedulerError(f"failed to spawn task {spec.name!r}: {exc}") from exc + + self.occupied_gpus.update(gpus) + logger.info( + "submitted %s pid=%d host=%s ports=%s gpus=%s", + spec.name, + pid, + alloc.host, + ports, + gpus, + ) + return TaskHandle( + host=alloc.host, + ports=ports, + gpu_devices=list(gpus), + ref={"pid": pid}, + ) + + def _allocate_gpus(self, n: int) -> list[int]: + if n <= 0: + return [] + free = [g for g in self.all_gpus if g not in self.occupied_gpus] + if len(free) < n: + raise SchedulerError( + f"need {n} GPUs but only {len(free)} free " + f"(total={self.all_gpus}, occupied={sorted(self.occupied_gpus)})" + ) + return free[:n] + + def _allocate_ports(self, n: int) -> list[int]: + ports = find_free_ports(n, exclude_ports=self._allocated_ports) + self._allocated_ports.update(ports) + return ports diff --git a/areal/v2/cli/inference/state.py b/areal/v2/cli/inference/state.py new file mode 100644 index 0000000000..21d0f0d212 --- /dev/null +++ b/areal/v2/cli/inference/state.py @@ -0,0 +1,348 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +from collections.abc import Iterable, Iterator +from contextlib import contextmanager +from dataclasses import asdict, dataclass, field +from pathlib import Path + +from areal.v2.cli.inference.scheduler import TaskHandle +from areal.v2.cli.process import pid_alive +from areal.v2.cli.state import ( + DEFAULT_SERVICE, + NamespacedStateStore, + ServiceStateBase, + SupportsComponentProbe, + atomic_write_json, +) +from areal.v2.cli.utils import file_lock + +INF_NAMESPACE = "inf" + + +class InferenceStateStore(NamespacedStateStore): + """Inference-specific extension of the scaffold ``NamespacedStateStore``. + + The inference service splits its state across two files per service — + the service-state (gateway / router handles, inherited from + scaffold) and the model-state (model registry, locked for + register / deregister). This class adds the model-state file paths + + lock + overrides ``recover_pids_from_raw_state`` to walk both + files when ``run --force`` falls back to raw-JSON PID recovery. + """ + + def __init__(self, namespace: str = INF_NAMESPACE) -> None: + super().__init__(namespace) + + # --- model-state file layout ---------------------------------------- + + def models_dir(self) -> Path: + """``$AREAL_HOME//models`` — created lazily.""" + + d = self.root() / "models" + d.mkdir(parents=True, exist_ok=True) + return d + + def models_state_path(self, service: str) -> Path: + return self.models_dir() / f"{service}.json" + + def models_lock_path(self, service: str) -> Path: + return self.models_dir() / f"{service}.lock" + + # --- locking -------------------------------------------------------- + + @contextmanager + def lock_model_state(self, service: str) -> Iterator[None]: + """Per-service flock on the model-registry file. Held across + register / deregister so concurrent CLI calls cannot race the + JSON read-modify-write.""" + + with file_lock(self.models_lock_path(service)): + yield + + # --- best-effort recovery (override) -------------------------------- + + def recover_pids_from_raw_state(self, service: str) -> list[int]: + """Walk BOTH service-state and model-state JSON for ``pid`` / + ``pids`` keys. Inference overrides scaffold's single-file walker + because the second (model-state) file is invisible to it.""" + + pids: list[int] = [] + pid_keys = {"pid", "pids"} + + def add(value) -> None: + if isinstance(value, int) and value > 0: + pids.append(value) + elif isinstance(value, list): + for item in value: + add(item) + + def walk(value) -> None: + if isinstance(value, dict): + for key, item in value.items(): + if key in pid_keys: + add(item) + else: + walk(item) + elif isinstance(value, list): + for item in value: + walk(item) + + for path in (self.service_state_path(service), self.models_state_path(service)): + if not path.exists(): + continue + with open(path) as f: + walk(json.load(f)) + + seen: set[int] = set() + unique: list[int] = [] + for pid in pids: + if pid not in seen: + seen.add(pid) + unique.append(pid) + return unique + + +# Module-level singleton — every dataclass / verb routes path resolution +# through this. Tests use ``AREAL_HOME`` to isolate; the store re-resolves +# on every call so the env override takes effect without rebuilding. +store = InferenceStateStore() + + +# ``DEFAULT_SERVICE`` is re-exported above so existing subcommand imports +# (`from ...inference.state import DEFAULT_SERVICE`) keep working. +__all__ = [ + "DEFAULT_SERVICE", + "INF_NAMESPACE", + "InferenceStateStore", + "ModelEntry", + "ModelReplica", + "ModelState", + "RuntimeState", + "ServiceState", + "store", +] + + +def _handle_from_dict(raw: dict) -> TaskHandle: + return TaskHandle( + host=raw["host"], + ports=list(raw.get("ports", [])), + gpu_devices=list(raw.get("gpu_devices", [])), + ref=dict(raw.get("ref", {})), + ) + + +@dataclass +class ModelReplica: + """One DP replica: data-proxy in front, worker behind. Field order + matches the request-flow direction used by stop / status output.""" + + data_proxy: TaskHandle + worker: TaskHandle + + +@dataclass +class ModelEntry: + backend: str = "" + replicas: list[ModelReplica] = field(default_factory=list) + + def all_workers(self) -> list[TaskHandle]: + return [r.worker for r in self.replicas] + + def all_data_proxies(self) -> list[TaskHandle]: + return [r.data_proxy for r in self.replicas] + + def gpu_devices(self) -> list[int]: + return [g for w in self.all_workers() for g in w.gpu_devices] + + @classmethod + def from_dict(cls, raw: dict) -> ModelEntry: + replicas = [ + ModelReplica( + data_proxy=_handle_from_dict(r["data_proxy"]), + worker=_handle_from_dict(r["worker"]), + ) + for r in raw.get("replicas", []) + ] + return cls(backend=raw.get("backend", ""), replicas=replicas) + + +@dataclass +class ServiceState: + service: str + backend: str + gateway_handle: TaskHandle + router_handle: TaskHandle + admin_api_key: str + started_at: float + launch_mode: str = "detached" + + @property + def gateway_pid(self) -> int: + return self.gateway_handle.pid + + @property + def gateway_url(self) -> str: + return self.gateway_handle.addr + + @property + def router_pid(self) -> int: + return self.router_handle.pid + + @property + def router_url(self) -> str: + return self.router_handle.addr + + def save(self) -> None: + atomic_write_json(store.service_state_path(self.service), asdict(self)) + store.set_current_service(self.service) + + @classmethod + def load(cls, service: str) -> ServiceState: + p = store.service_state_path(service) + if not p.exists(): + raise FileNotFoundError(f"No service state at {p}") + with open(p) as f: + raw = json.load(f) + return cls( + service=raw.get("service", service), + backend=raw["backend"], + gateway_handle=_handle_from_dict(raw["gateway_handle"]), + router_handle=_handle_from_dict(raw["router_handle"]), + admin_api_key=raw["admin_api_key"], + started_at=float(raw["started_at"]), + launch_mode=raw.get("launch_mode", "detached"), + ) + + @classmethod + def remove(cls, service: str) -> None: + p = store.service_state_path(service) + if p.exists(): + p.unlink() + store.clear_current_service(service) + + +@dataclass +class ModelState: + service: str + models: dict[str, ModelEntry] = field(default_factory=dict) + + def save(self) -> None: + atomic_write_json(store.models_state_path(self.service), asdict(self)) + + @classmethod + def load(cls, service: str) -> ModelState: + p = store.models_state_path(service) + if not p.exists(): + return cls(service=service) + with open(p) as f: + raw = json.load(f) + models = { + name: ModelEntry.from_dict(entry) + for name, entry in (raw.pop("models", None) or {}).items() + } + return cls(service=raw.get("service", service), models=models) + + @classmethod + def remove(cls, service: str) -> None: + p = store.models_state_path(service) + if p.exists(): + p.unlink() + + def all_workers(self) -> list[TaskHandle]: + return [h for entry in self.models.values() for h in entry.all_workers()] + + def all_data_proxies(self) -> list[TaskHandle]: + return [h for entry in self.models.values() for h in entry.all_data_proxies()] + + def occupied_gpus(self) -> set[int]: + return {g for entry in self.models.values() for g in entry.gpu_devices()} + + +@dataclass +class RuntimeState(ServiceStateBase): + """Composite of ServiceState (gateway/router) and ModelState (model + registry). The single object the CLI verbs operate on; the lifecycle + contract is implemented against ``RuntimeState`` so scaffold helpers + enumerate every component in one pass.""" + + service_state: ServiceState + model_state: ModelState + + @property + def service(self) -> str: + return self.service_state.service + + @property + def backend(self) -> str: + return self.service_state.backend + + @property + def gateway_handle(self) -> TaskHandle: + return self.service_state.gateway_handle + + @property + def router_handle(self) -> TaskHandle: + return self.service_state.router_handle + + @property + def gateway_pid(self) -> int: + return self.service_state.gateway_pid + + @property + def gateway_url(self) -> str: + return self.service_state.gateway_url + + @property + def router_pid(self) -> int: + return self.service_state.router_pid + + @property + def router_url(self) -> str: + return self.service_state.router_url + + @property + def admin_api_key(self) -> str: + return self.service_state.admin_api_key + + @property + def started_at(self) -> float: + return self.service_state.started_at + + @property + def launch_mode(self) -> str: + return self.service_state.launch_mode + + @property + def models(self) -> dict[str, ModelEntry]: + return self.model_state.models + + def save(self) -> None: + self.service_state.save() + self.model_state.save() + + @classmethod + def load(cls, service: str) -> RuntimeState: + service_state = ServiceState.load(service) + return cls( + service_state=service_state, + model_state=ModelState.load(service_state.service), + ) + + # --- ServiceStateBase contract ---------------------------------------- + + def gateway_alive(self) -> bool: + # Non-local backends will need their own liveness probe — PID-based + # check returns False for them (no pid in handle.ref). + return pid_alive(self.gateway_pid) + + def components(self) -> Iterable[tuple[str, SupportsComponentProbe]]: + yield "gateway", self.gateway_handle + yield "router", self.router_handle + for name, entry in self.model_state.models.items(): + for i, replica in enumerate(entry.replicas): + yield f"data_proxy[{name}/{i}]", replica.data_proxy + yield f"worker[{name}/{i}]", replica.worker diff --git a/areal/v2/cli/lifecycle.py b/areal/v2/cli/lifecycle.py new file mode 100644 index 0000000000..5f6b8985fc --- /dev/null +++ b/areal/v2/cli/lifecycle.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Service lifecycle helpers shared across subcommand CLIs. + +A ``ServiceLifecycle`` bundles a CLI's namespace + state class + stop +command name, then exposes the three "is this service running" +predicates (``running_state`` / ``load_running_state`` / +``refuse_if_running``) plus ``force_replace_slot`` for ``run --force``. + +Subcommand CLIs subclass to inject their state class and may override +``gateway_alive`` if they don't want the default +``state.gateway_alive()`` delegation. +""" + +from __future__ import annotations + +from pathlib import Path + +import click + +from areal.v2.cli.process import kill_pids +from areal.v2.cli.state import DEFAULT_SERVICE, NamespacedStateStore + + +class ServiceLifecycle: + namespace: str = "" + state_class: type = type(None) + stop_command: str = "" + default_service: str = DEFAULT_SERVICE + + def __init__( + self, + *, + namespace: str | None = None, + state_class: type | None = None, + stop_command: str | None = None, + ) -> None: + if namespace is not None: + self.namespace = namespace + if state_class is not None: + self.state_class = state_class + if stop_command is not None: + self.stop_command = stop_command + if not self.namespace: + raise ValueError("ServiceLifecycle requires a namespace") + if not self.stop_command: + raise ValueError("ServiceLifecycle requires a stop_command") + self.store = NamespacedStateStore(self.namespace) + + # ------------------------------------------------------------------ + # Hooks subclasses may override + + def gateway_alive(self, state) -> bool: + """Default delegates to ``state.gateway_alive()`` (the + ServiceStateBase contract). Subclasses may override if their + state object can't or shouldn't implement that method itself.""" + + return state.gateway_alive() + + def state_path(self, service: str) -> Path: + return self.store.service_state_path(service) + + def resolve_service_name(self, explicit: str | None) -> str: + return self.store.resolve_service_name(explicit, fallback=self.default_service) + + def list_services(self) -> list[str]: + return self.store.list_service_names() + + def load_state(self, service: str): + return self.state_class.load(service) + + # ------------------------------------------------------------------ + # Public API + + def running_state(self, service: str | None = None): + """Return the loaded state iff it exists, parses, and reports + the gateway alive. Returns None otherwise — never raises.""" + + name = self.resolve_service_name(service) + if not self.state_path(name).exists(): + return None + try: + state = self.load_state(name) + except Exception: + return None + if not self.gateway_alive(state): + return None + return state + + def load_running_state(self, service: str | None = None): + """Like ``running_state`` but raises ``ClickException`` on + every failure — for command bodies that need a state to proceed.""" + + name = self.resolve_service_name(service) + if not self.state_path(name).exists(): + raise click.ClickException(f"service {name!r} is not running") + try: + state = self.load_state(name) + except Exception as exc: + raise click.ClickException(f"failed to load state: {exc}") from exc + if not self.gateway_alive(state): + raise click.ClickException(f"service {name!r} gateway is not alive") + return state + + def refuse_if_running(self, service: str | None = None) -> None: + """Used by ``run`` to refuse double-start. No-op if no running + service exists; raises ``ClickException`` if one does.""" + + state = self.running_state(service) + if state is None: + return + raise click.ClickException( + f"service {state.service!r} already running. " + f"Run `{self.stop_command} --service {state.service}` first." + ) + + def force_replace_slot(self, service: str, *, grace_s: float = 5.0) -> None: + """``run --force`` path: tear down any existing children for + ``service`` (loading the state for an orderly kill, falling back + to the raw-JSON PID walk if the state file is corrupted), then + unlink the state file so the next spawn starts clean. + + Caller is still responsible for the subsequent fresh spawn. + """ + + path = self.state_path(service) + if not path.exists(): + return + pids: list[int] = [] + try: + state = self.load_state(service) + pids = self._collect_pids(state) + except Exception: + try: + pids = self.store.recover_pids_from_raw_state(service) + except Exception: + pids = [] + if pids: + kill_pids(pids, grace_s=grace_s) + path.unlink(missing_ok=True) + self.store.clear_current_service(service) + + def _collect_pids(self, state) -> list[int]: + """Default implementation pulls ``.pid`` off every component + ``state.components()`` yields. Subclasses with extra state files + (e.g. inf's ``models/.json``) may override to extend the + list.""" + + return [pid for _, h in state.components() if (pid := h.pid) > 0] diff --git a/areal/v2/cli/main.py b/areal/v2/cli/main.py new file mode 100644 index 0000000000..4d723a724d --- /dev/null +++ b/areal/v2/cli/main.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from areal.v2.cli.cli import cli + +if __name__ == "__main__": + cli() diff --git a/areal/v2/cli/process.py b/areal/v2/cli/process.py new file mode 100644 index 0000000000..9df92fda06 --- /dev/null +++ b/areal/v2/cli/process.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Local-only process primitives for service-style CLIs. + +These helpers operate on local PIDs and process trees. Remote backends +(K8s / Slurm) should not consume them directly — route teardown through +``areal.v2.cli.scheduler.terminate`` instead, which dispatches +per backend. + +For port allocation use ``areal.utils.network.find_free_ports``; it +draws from a non-ephemeral range and supports excluding already-handed- +out ports, avoiding the TOCTOU collisions a naive ``bind(0)`` walks +into. +""" + +from __future__ import annotations + +import os +import subprocess +from pathlib import Path + +from areal.infra.utils.proc import kill_process_tree + + +def pid_alive(pid: int) -> bool: + if pid <= 0: + return False + # If ``pid`` is one of our own children we may catch it mid-zombie: + # ``os.kill(pid, 0)`` still succeeds for zombies, so an unreaped child + # would look alive forever and kill_pids would burn its full grace + # window before sending a redundant SIGKILL. Reap it first via WNOHANG + # so a zombie is reported dead immediately. + try: + reaped, _ = os.waitpid(pid, os.WNOHANG) + if reaped == pid: + return False + except (ChildProcessError, OSError): + # Not our child, or already reaped — fall through to the kill probe. + pass + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + +def spawn_process( + cmd: list[str], log_file: Path, env: dict[str, str] | None = None +) -> int: + """Spawn a detached subprocess that survives parent exit. + + ``start_new_session=True`` puts the child in its own session so the + parent receiving SIGHUP (terminal close) does not propagate to the + child. stdout / stderr are appended to ``log_file``. Extra env vars + in ``env`` are merged on top of the parent environment. + """ + + log_file.parent.mkdir(parents=True, exist_ok=True) + final_env = os.environ.copy() + final_env.setdefault("PYTHONUNBUFFERED", "1") + if env: + final_env.update(env) + # Popen dup()s the fd for the child, so closing our copy here does not + # affect the child's stdout/stderr — and it stops us from leaking a fd + # in the parent every time spawn_process is called. + with open(log_file, "ab", buffering=0) as log_handle: + proc = subprocess.Popen( + cmd, + stdin=subprocess.DEVNULL, + stdout=log_handle, + stderr=subprocess.STDOUT, + start_new_session=True, + env=final_env, + ) + return proc.pid + + +def kill_pids(pids: list[int], grace_s: float) -> None: + """SIGTERM → wait ``grace_s`` → SIGKILL across local pids and their + descendant process trees. + + Delegates per-pid to ``areal.infra.utils.proc.kill_process_tree``, + which walks ``psutil.Process(pid).children(recursive=True)`` so + grandchildren that escaped the original process group are still + caught. Local-only — for remote backends use + ``areal.v2.cli.scheduler.terminate``. + """ + + for pid in pids: + if pid > 0: + kill_process_tree(pid, timeout=int(grace_s), graceful=True) diff --git a/areal/v2/cli/scheduler.py b/areal/v2/cli/scheduler.py new file mode 100644 index 0000000000..750d828f63 --- /dev/null +++ b/areal/v2/cli/scheduler.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Unified termination dispatch for service-style CLIs. + +Provides a stateless top-level ``terminate`` that takes a TaskHandle's +``ref`` dict plus the backend tag, and routes to the matching kill +path. Currently only the local backend is implemented; future backends +extend this module with their own branches. +""" + +from __future__ import annotations + +from areal.v2.cli.process import kill_pids + + +def terminate( + ref: dict, + *, + backend: str = "local", + grace_s: float = 10.0, +) -> None: + """Terminate the task identified by ``ref`` using the given backend. + + For ``backend="local"`` ``ref`` must contain a ``"pid"`` key — the + local PID and its descendant process tree are SIGTERM'd, given + ``grace_s`` seconds to exit, then SIGKILL'd. + """ + + if backend == "local": + pid = int(ref.get("pid", 0) or 0) + if pid > 0: + kill_pids([pid], grace_s=grace_s) + return + + raise ValueError(f"unknown scheduler backend: {backend!r}") diff --git a/areal/v2/cli/state.py b/areal/v2/cli/state.py new file mode 100644 index 0000000000..251946760b --- /dev/null +++ b/areal/v2/cli/state.py @@ -0,0 +1,280 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""On-disk state primitives shared across service-style CLIs. + +Three layers live here: + +1. **Global helpers** — ``areal_home`` (the user's CLI root) and + ``atomic_write_json`` (a generic write primitive). Both are + stateless; no namespace involvement. + +2. **NamespacedStateStore** — every subcommand CLI binds a namespace + (``inf``, ``agent``, ``train``, …) and lives under + ``$AREAL_HOME//``. The store class collects the path + resolution / pointer file / orphan-recovery operations so callers + construct one instance per namespace and use methods on it rather + than threading the namespace string through every call site. + +3. **Contract types** — ``SupportsComponentProbe`` (Protocol) and + ``ServiceStateBase`` (ABC). Subcommand CLIs implement their own + ServiceState dataclass that satisfies the protocol / base class, and + in return get to plug into scaffold's ``ServiceLifecycle`` / + ``StatusReporter`` / etc. without further glue. +""" + +from __future__ import annotations + +import json +import os +import tempfile +from abc import ABC, abstractmethod +from collections.abc import Iterable +from pathlib import Path +from typing import Any, Protocol, runtime_checkable + +DEFAULT_SERVICE = "default" + + +# --------------------------------------------------------------------------- +# Global helpers +# --------------------------------------------------------------------------- + + +def areal_home() -> Path: + """Return the AReaL CLI home directory. + + Resolves ``$AREAL_HOME`` if set, otherwise ``~/.areal``. Created on + first access so callers can mkdir-then-write subdirs without an + explicit setup step. + """ + + env = os.environ.get("AREAL_HOME") + root = Path(env).expanduser() if env else Path.home() / ".areal" + root.mkdir(parents=True, exist_ok=True) + return root + + +def atomic_write_json(path: Path, data: Any, *, indent: int = 2) -> None: + """Atomically write ``data`` as JSON to ``path``. + + Writes to a unique tempfile in ``path``'s directory, fsync()s it to + disk, then renames into place. ``NamedTemporaryFile(delete=False)`` + gives us a fresh name per call so concurrent writers do not stomp on + each other's tempfiles, and the tempfile is unlinked on serialization + or rename failure so half-written state never lingers on disk. + """ + + path.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + "w", dir=path.parent, delete=False, suffix=".tmp" + ) as f: + tmp_path = Path(f.name) + try: + json.dump(data, f, indent=indent, default=str) + f.write("\n") + f.flush() + os.fsync(f.fileno()) + except Exception: + tmp_path.unlink(missing_ok=True) + raise + try: + os.replace(tmp_path, path) + except Exception: + tmp_path.unlink(missing_ok=True) + raise + + +# --------------------------------------------------------------------------- +# Namespace-bound state store +# --------------------------------------------------------------------------- + + +class NamespacedStateStore: + """All ``$AREAL_HOME//...`` path resolution + pointer + file management for one subcommand CLI. + + Construct one per CLI (e.g. ``store = NamespacedStateStore("inf")``) + and use the methods on it. Paths are re-resolved on every call so + tests that override the home directory via ``AREAL_HOME`` work + without rebuilding the store. + """ + + def __init__(self, namespace: str) -> None: + if not namespace: + raise ValueError("NamespacedStateStore requires a non-empty namespace") + self.namespace = namespace + + # --- directory roots ------------------------------------------------- + + def root(self) -> Path: + d = areal_home() / self.namespace + d.mkdir(parents=True, exist_ok=True) + return d + + def services_dir(self) -> Path: + d = self.root() / "services" + d.mkdir(parents=True, exist_ok=True) + return d + + def logs_root(self) -> Path: + d = self.root() / "logs" + d.mkdir(parents=True, exist_ok=True) + return d + + def logs_dir(self, service: str) -> Path: + d = self.logs_root() / service + d.mkdir(parents=True, exist_ok=True) + return d + + # --- file paths ------------------------------------------------------ + + def service_state_path(self, service: str) -> Path: + return self.services_dir() / f"{service}.json" + + def service_lock_path(self, service: str) -> Path: + return self.services_dir() / f"{service}.lock" + + def current_service_path(self) -> Path: + return self.root() / "current-service" + + def config_path(self) -> Path: + return self.root() / "config.toml" + + # --- pointer file ---------------------------------------------------- + + def list_service_names(self) -> list[str]: + return sorted(p.stem for p in self.services_dir().glob("*.json")) + + def set_current_service(self, service: str) -> None: + self.current_service_path().write_text(service + "\n") + + def clear_current_service(self, service: str) -> None: + path = self.current_service_path() + if path.exists() and path.read_text().strip() == service: + path.unlink() + + def resolve_service_name( + self, + explicit: str | None = None, + *, + fallback: str = DEFAULT_SERVICE, + ) -> str: + """Resolve the active service name for a CLI call. + + Order: ``--service`` flag > current-service pointer file > the + single running service (if exactly one) > ``fallback``. + """ + + if explicit: + return explicit + pointer = self.current_service_path() + if pointer.exists(): + value = pointer.read_text().strip() + if value: + return value + running = self.list_service_names() + if len(running) == 1: + return running[0] + return fallback + + # --- best-effort orphan PID recovery --------------------------------- + + def recover_pids_from_raw_state(self, service: str) -> list[int]: + """Walk the service-state file for ``service`` and pull any + ``pid`` / ``pids`` numbers. + + Used by ``run --force`` to clean up children when the dataclass + parse fails (state file from an older / corrupted schema). + Subclasses with extra state files (e.g. inf's + ``models/.json``) override to walk the additional files. + """ + + pids: list[int] = [] + pid_keys = {"pid", "pids"} + + def add(value) -> None: + if isinstance(value, int) and value > 0: + pids.append(value) + elif isinstance(value, list): + for item in value: + add(item) + + def walk(value) -> None: + if isinstance(value, dict): + for key, item in value.items(): + if key in pid_keys: + add(item) + else: + walk(item) + elif isinstance(value, list): + for item in value: + walk(item) + + path = self.service_state_path(service) + if path.exists(): + with open(path) as f: + walk(json.load(f)) + + seen: set[int] = set() + unique: list[int] = [] + for pid in pids: + if pid not in seen: + seen.add(pid) + unique.append(pid) + return unique + + +# --------------------------------------------------------------------------- +# Contract types +# --------------------------------------------------------------------------- + + +@runtime_checkable +class SupportsComponentProbe(Protocol): + """Minimum surface a subcommand's component handle must expose so + scaffold helpers can probe it and identify it. + + Each subcommand CLI's handle type — inf's ``TaskHandle``, agent's + ``ProcessState``, etc. — already provides ``addr`` (HTTP base for + ``/health`` probes) and ``pid`` (for local liveness / kill paths). + No inheritance needed: structural duck typing. + """ + + @property + def addr(self) -> str: ... + + @property + def pid(self) -> int: ... + + +class ServiceStateBase(ABC): + """Abstract base every subcommand CLI's ServiceState should satisfy. + + The base nails down the universal fields and the two methods + (``gateway_alive`` / ``components``) that scaffold's lifecycle and + status reporters rely on. Subclasses are free to add backend-specific + fields (engine handles, model registries, agent pair configs, etc.). + """ + + service: str + admin_api_key: str + launch_mode: str + started_at: float + + @abstractmethod + def gateway_alive(self) -> bool: + """Return True iff the service's central entry point is reachable. + + ``ServiceLifecycle`` uses this to decide "running" — every CLI + must define what alive means for its own architecture (local PID + alive / k8s pod ready / slurm job state). + """ + + @abstractmethod + def components(self) -> Iterable[tuple[str, SupportsComponentProbe]]: + """Yield ``(label, handle)`` for every component of this service. + + Used by ``StatusReporter`` to enumerate rows; the order is the + order rows appear in the table. Labels are display-only strings + (e.g. ``"gateway"``, ``"worker[qwen/0]"``). + """ diff --git a/areal/v2/cli/status.py b/areal/v2/cli/status.py new file mode 100644 index 0000000000..f3389168d6 --- /dev/null +++ b/areal/v2/cli/status.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Status drill-down table with parallel /health probing. + +A ``StatusReporter`` is constructed with the components to inspect plus +a list of ``ColumnSpec`` describing the columns to render. It probes +every component's ``/health`` endpoint concurrently and produces either +a text table or a JSON payload. + +Both the column set and the per-column value extractors are caller- +supplied — subcommand CLIs decide what to show. Scaffold provides only +the orchestration (parallel probe + table rendering). +""" + +from __future__ import annotations + +import urllib.error +import urllib.request +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass + +import click + +from areal.v2.cli.state import SupportsComponentProbe + + +@dataclass +class ColumnSpec: + """One column in the status table. + + ``value`` receives the row's ``(label, handle, alive)`` tuple and + returns the rendered string. Returning ``"-"`` for "not applicable" + is the conventional placeholder. + """ + + header: str + value: Callable[[str, SupportsComponentProbe, bool], str] + + +class StatusReporter: + def __init__( + self, + components: list[tuple[str, SupportsComponentProbe]], + columns: list[ColumnSpec], + *, + probe_timeout: float = 3.0, + max_workers: int = 8, + ) -> None: + self.components = list(components) + self.columns = list(columns) + self.probe_timeout = probe_timeout + self.max_workers = max_workers + + def probe_all(self) -> list[bool]: + """Probe every component's ``/health`` endpoint concurrently. + + Returns a list of bools aligned with ``self.components``. A + non-5xx response within ``probe_timeout`` counts as alive; any + network error or 5xx counts as not alive. + """ + + if not self.components: + return [] + addrs = [handle.addr or "" for _, handle in self.components] + + def probe(addr: str) -> bool: + if not addr: + return False + try: + with urllib.request.urlopen( + f"{addr.rstrip('/')}/health", timeout=self.probe_timeout + ) as resp: + return resp.status < 500 + except ( + urllib.error.URLError, + ConnectionError, + TimeoutError, + OSError, + ): + return False + + with ThreadPoolExecutor( + max_workers=max(1, min(self.max_workers, len(addrs))) + ) as pool: + return list(pool.map(probe, addrs)) + + def render_rows(self, alive_flags: list[bool] | None = None) -> list[list[str]]: + """Build per-row string lists using each column's value extractor.""" + + if alive_flags is None: + alive_flags = self.probe_all() + rows: list[list[str]] = [] + for (label, handle), alive in zip(self.components, alive_flags, strict=True): + rows.append([col.value(label, handle, alive) for col in self.columns]) + return rows + + def print_table( + self, + rows: list[list[str]] | None = None, + *, + header_line: str | None = None, + ) -> None: + """Print a left-aligned table with column widths derived from + the longest cell in each column. Optional ``header_line`` is + echoed before the table (e.g. ``"service: foo backend: local"``). + """ + + if rows is None: + rows = self.render_rows() + if header_line: + click.echo(header_line) + click.echo() + headers = [c.header for c in self.columns] + widths = [ + max(len(h), *(len(r[i]) for r in rows)) for i, h in enumerate(headers) + ] + fmt = " ".join(f"{{:<{w}}}" for w in widths) + click.echo(fmt.format(*headers)) + for row in rows: + click.echo(fmt.format(*row)) + + def json_snapshot(self, alive_flags: list[bool] | None = None) -> list[dict]: + """Return the table content as a list of dicts (one per row) + keyed by column header — convenient for ``--json`` output.""" + + if alive_flags is None: + alive_flags = self.probe_all() + snapshot = [] + for (label, handle), alive in zip(self.components, alive_flags, strict=True): + snapshot.append( + {col.header: col.value(label, handle, alive) for col in self.columns} + ) + return snapshot diff --git a/areal/v2/cli/training/__init__.py b/areal/v2/cli/training/__init__.py new file mode 100644 index 0000000000..1d392ca050 --- /dev/null +++ b/areal/v2/cli/training/__init__.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import click + +from areal.v2.cli.training.commands.run import run_cmd + + +@click.group(help="Run AReaL training experiments.") +def train() -> None: + pass + + +train.add_command(run_cmd) diff --git a/areal/v2/cli/training/cli_guide.md b/areal/v2/cli/training/cli_guide.md new file mode 100644 index 0000000000..3455da8704 --- /dev/null +++ b/areal/v2/cli/training/cli_guide.md @@ -0,0 +1,139 @@ +# AReaL Training Service CLI + +`areal train` is the training subcommand group under the top-level `areal` CLI. It wires +up an AReaL training driver function and an experiment config file into a single +command. It does not manage the training process lifecycle (unlike `areal inf`, which +maintains service state); it only "finds the driver, loads the config, and passes +hydra-style overrides through." + +## Basic concepts + +The minimum execution unit of a training job is a **driver function** —typically a +`main(args: list[str])` in some script under `examples/`. The CLI does exactly three +things: + +1. Resolve the driver from `module.path:func` +1. Resolve `--config ` to an absolute path and prepend it to argv +1. Append every trailing argument unchanged to argv (typically hydra overrides) + +The driver function's return value is used as the exit code if it returns `int`; +anything else (including `None`) is treated as 0. + +## Usage + +```bash +areal train run \ + --config \ + --driver : \ + [ ...] +``` + +| flag / arg | required | description | +| ------------------------ | -------- | ----------------------------------------------------------------------------- | +| `--config` | yes | Experiment YAML path; the file must exist (the CLI runs an `exists` check) | +| `--driver` | yes | Driver entry point in `module.path:func` form (colon-separated) | +| trailing positional args | no | Forwarded to the driver verbatim; typically hydra-style `key=value` overrides | + +Note that `run_cmd` is configured with +`context_settings={"ignore_unknown_options": True}` — this means trailing positional +args **can include `--xxx` flags**; the CLI does not try to parse them and forwards them +as-is to the driver. + +## Examples + +Run GSM8K GRPO (the most common baseline): + +```bash +areal train run \ + --config examples/math/gsm8k_grpo.yaml \ + --driver examples.math.gsm8k_rl:main \ + experiment_name=gsm8k_grpo_test \ + trial_name=t1 +``` + +Run SFT: + +```bash +areal train run \ + --config examples/math/gsm8k_sft.yaml \ + --driver examples.math.gsm8k_sft:main +``` + +## Driver function conventions + +The CLI calls the driver with a single argument `argv: list[str]`, so the driver must +look like: + +```python +def main(args: list[str]) -> int | None: + config, _ = load_expr_config(args, GRPOConfig) # or any other *Config dataclass + ... + return 0 +``` + +`load_expr_config` lives in `areal.api.cli_args` and consumes `args` itself: it +recognises the YAML path after `--config`, and treats every remaining `key=value` as a +hydra override merged into the config dataclass. In other words, hydra parsing is **done +by the driver**, not the CLI. + +Minimum template for writing a new driver: + +```python +from areal.api.cli_args import GRPOConfig, load_expr_config +from areal import PPOTrainer + +def main(args): + config, _ = load_expr_config(args, GRPOConfig) + with PPOTrainer(config, train_dataset=..., valid_dataset=...) as trainer: + trainer.train(workflow="...", workflow_kwargs={...}) + return 0 +``` + +## Hydra overrides + +Any driver that uses `load_expr_config` to parse args supports hydra-style overrides. +Common override targets: + +```bash +# experiment / trial naming +experiment_name=my_run trial_name=t1 + +# cluster size +cluster.n_nodes=4 cluster.n_gpus_per_node=8 + +# training hyperparameters +actor.optimizer.lr=5e-6 +total_train_epochs=20 + +# rollout backend +rollout.backend=sglang:d2p1t2 +rollout.max_concurrent_rollouts=128 + +# datasets +train_dataset.batch_size=256 +``` + +The CLI does not validate whether these keys are legal; unknown fields will be reported +by hydra when the driver loads the config. + +## Exit codes + +| Scenario | exit code | +| ------------------------------------------------------ | ----------------------------------------- | +| Driver returns `int` | Returned value used directly as exit code | +| Driver returns `None` / other | 0 | +| `--driver` does not contain `:` | UsageError (click default 2) | +| Module referenced by `--driver` cannot be imported | ClickException (1) | +| Function referenced by `--driver` is not on the module | ClickException (1) | +| `--config` path does not exist | Caught by click `exists=True` (2) | + +Exceptions raised inside the driver **are not caught by the CLI** — the default Python +behaviour applies (traceback printed, process exits). + +## Not implemented yet + +`areal train` currently only implements `run`. The following are reasonable future +extensions but are **not** in this version: + +- `areal train ps` / `status` / `stop` — lifecycle management for training jobs + (requires a training service state concept first) diff --git a/areal/experimental/inference_service/vllm/__init__.py b/areal/v2/cli/training/commands/__init__.py similarity index 100% rename from areal/experimental/inference_service/vllm/__init__.py rename to areal/v2/cli/training/commands/__init__.py diff --git a/areal/v2/cli/training/commands/run.py b/areal/v2/cli/training/commands/run.py new file mode 100644 index 0000000000..104e5df72b --- /dev/null +++ b/areal/v2/cli/training/commands/run.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import importlib +from pathlib import Path + +import click + + +@click.command( + name="run", + help="Invoke a training driver with the given config and hydra overrides.", + context_settings={"ignore_unknown_options": True}, +) +@click.option( + "--config", + "config_path", + type=click.Path(exists=True, dir_okay=False, path_type=Path), + required=True, + help="Path to the experiment yaml.", +) +@click.option( + "--driver", + "driver_spec", + required=True, + help="Driver entry point as 'module.path:func', e.g. examples.math.gsm8k_rl:main.", +) +@click.argument("overrides", nargs=-1, type=click.UNPROCESSED) +def run_cmd(config_path: Path, driver_spec: str, overrides: tuple[str, ...]) -> None: + raise SystemExit(do_run(config_path, driver_spec, list(overrides)) or 0) + + +def do_run(config_path: Path, driver_spec: str, overrides: list[str]) -> int: + if ":" not in driver_spec: + raise click.UsageError( + f"--driver must be in 'module.path:func' form, got: {driver_spec!r}" + ) + mod_path, func_name = driver_spec.split(":", 1) + try: + module = importlib.import_module(mod_path) + except ImportError as e: + raise click.ClickException( + f"failed to import driver module {mod_path!r}: {e}" + ) from e + try: + fn = getattr(module, func_name) + except AttributeError as e: + raise click.ClickException( + f"module {mod_path!r} has no attribute {func_name!r}" + ) from e + + argv = ["--config", str(config_path.resolve()), *overrides] + result = fn(argv) + if isinstance(result, int): + return result + return 0 diff --git a/areal/v2/cli/utils.py b/areal/v2/cli/utils.py new file mode 100644 index 0000000000..47a0a45604 --- /dev/null +++ b/areal/v2/cli/utils.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Stateless utility helpers shared across subcommand CLIs. + +Four groups live here: + +- :func:`file_lock` — per-service mutual exclusion via POSIX flock. +- :func:`wait_http_health` / :func:`wait_client_health` — block until a + newly-spawned component reports ready. +- :func:`json_or_table` — ``--json`` vs table dual-output dispatch. +- :func:`register_cli_logger` — register a CLI logger name + color in + the global color table and return the configured Logger. +""" + +from __future__ import annotations + +import fcntl +import json +import logging +import time +import urllib.error +import urllib.request +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Protocol + +import click + +from areal.utils.logging import LOGGER_COLORS_EXACT, getLogger +from areal.v2.cli.process import pid_alive + +# --------------------------------------------------------------------------- +# File locking +# --------------------------------------------------------------------------- + + +@contextmanager +def file_lock(path: Path) -> Iterator[None]: + """Hold an exclusive flock on ``path`` for the duration of the ``with`` block. + + The lock file is created if missing; concurrent waiters block until + the current holder exits the ``with`` block. Used by subcommand CLIs + to serialize concurrent ``register`` / ``deregister`` / ``stop`` + calls against the same service's state file. + """ + + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "a+") as fp: + fcntl.flock(fp.fileno(), fcntl.LOCK_EX) + try: + yield + finally: + fcntl.flock(fp.fileno(), fcntl.LOCK_UN) + + +# --------------------------------------------------------------------------- +# Health polling +# --------------------------------------------------------------------------- + + +class _HealthCheckable(Protocol): + def health(self, *, timeout: float) -> object: ... + + +def wait_http_health( + url: str, + *, + pid: int, + timeout: float, + label: str, + poll_interval: float = 0.5, + request_timeout: float = 2.0, +) -> None: + """Poll ``GET /health`` until it returns < 5xx. + + Raises ``ClickException`` if the deadline elapses, or if the spawned + PID dies during startup (cheap early-fail signal for the local + case). Pass ``pid=0`` for non-local backends to skip the PID check. + """ + + deadline = time.time() + timeout + last_err: Exception | None = None + while time.time() < deadline: + if pid > 0 and not pid_alive(pid): + raise click.ClickException(f"{label} subprocess died during startup") + try: + with urllib.request.urlopen( + f"{url.rstrip('/')}/health", timeout=request_timeout + ) as resp: + if resp.status < 500: + return + except (urllib.error.URLError, ConnectionError, TimeoutError, OSError) as exc: + last_err = exc + time.sleep(poll_interval) + raise click.ClickException(f"{label} did not become healthy: {last_err}") + + +def wait_client_health( + client: _HealthCheckable, + *, + timeout: float, + label: str, + poll_interval: float = 0.3, + request_timeout: float = 1.5, +) -> None: + """Poll ``client.health(timeout=...)`` until it returns without raising. + + Any exception from ``client.health`` is treated as "not ready yet" + and the loop sleeps until the deadline. + """ + + deadline = time.time() + timeout + last_err: Exception | None = None + while time.time() < deadline: + try: + client.health(timeout=request_timeout) + return + except Exception as exc: + last_err = exc + time.sleep(poll_interval) + raise click.ClickException( + f"{label} did not become healthy within {timeout:.0f}s (last error: {last_err})" + ) + + +# --------------------------------------------------------------------------- +# Output dispatch +# --------------------------------------------------------------------------- + + +def json_or_table( + payload: Any, + *, + as_json: bool, + table_renderer: Callable[[Any], None], + indent: int = 2, +) -> None: + """Emit ``payload`` as JSON or hand it to a table renderer. + + Most subcommand ``ps`` / ``status`` / ``models`` verbs follow the + same pattern: ``--json`` dumps the raw payload, otherwise a table + is rendered. Funneling both through this helper keeps the branching + out of the command body. + """ + + if as_json: + click.echo(json.dumps(payload, indent=indent, default=str)) + return + table_renderer(payload) + + +# --------------------------------------------------------------------------- +# Logger registration +# --------------------------------------------------------------------------- + + +def register_cli_logger(name: str, color: str = "blue") -> logging.Logger: + """Register ``name`` with ``color`` in ``areal.utils.logging``'s + color table and return the configured Logger. + + Service-style CLIs conventionally use ``"blue"`` (infrastructure / + scheduler category). See ``areal/utils/logging.py`` for the palette + if a different color fits. + """ + + LOGGER_COLORS_EXACT[name] = color + return getLogger(name) diff --git a/areal/v2/cli/watcher.py b/areal/v2/cli/watcher.py new file mode 100644 index 0000000000..806717a6b5 --- /dev/null +++ b/areal/v2/cli/watcher.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Foreground watcher with signal-aware teardown. + +Subcommand CLIs use this in their ``run`` (non-detach) path to block +until the service exits, with a clear distinction between: + + - SIGINT / SIGTERM → user-requested stop → run the teardown callback + (kill workers, remove state file) → return 0 + - SIGHUP → terminal disconnect → exit quietly without teardown; + workers spawned with ``start_new_session=True`` survive + +Subclasses or constructor kwargs can swap the SIGHUP semantics for +``"teardown"`` if a particular CLI wants the older "all signals tear +down" behavior. +""" + +from __future__ import annotations + +import signal +import time +from collections.abc import Callable +from typing import Literal + +Disposition = Literal["teardown", "detach"] + + +class ForegroundWatcher: + """Block the calling thread until the service exits. + + Construct with an ``is_alive`` callback (typically + ``lambda: lifecycle.gateway_alive(state)``) and a ``teardown`` + callback (kills workers + removes state file). Call ``watch()`` — + it returns the desired CLI exit code. + """ + + default_signal_dispositions: dict[int, Disposition] = { + signal.SIGTERM: "teardown", + } + if hasattr(signal, "SIGHUP"): + default_signal_dispositions[signal.SIGHUP] = "detach" + + def __init__( + self, + *, + is_alive: Callable[[], bool], + teardown: Callable[[], None], + idle_poll: float = 1.0, + service_name: str = "", + signal_dispositions: dict[int, Disposition] | None = None, + ) -> None: + self.is_alive = is_alive + self.teardown = teardown + self.idle_poll = idle_poll + self.service_name = service_name + self.dispositions = ( + dict(signal_dispositions) + if signal_dispositions is not None + else dict(self.default_signal_dispositions) + ) + + def watch(self) -> int: + previous = self._install_handlers() + try: + while self.is_alive(): + time.sleep(self.idle_poll) + except KeyboardInterrupt: + # SIGINT (Ctrl+C) plus any SIGTERM-dispatched teardown lands here. + self.teardown() + return 0 + except SystemExit as exc: + # SIGHUP-dispatched detach path lands here. + return exc.code if isinstance(exc.code, int) else 0 + except BaseException: + self.teardown() + raise + finally: + self._restore_handlers(previous) + # is_alive() flipped to False on its own — gateway died externally. + # We do not call teardown; the children are already gone, but the + # caller may want to clean up state (handled outside this class). + return 0 + + # ------------------------------------------------------------------ + # Signal plumbing + + def _install_handlers(self) -> dict[int, object]: + previous: dict[int, object] = {} + for sig, disposition in self.dispositions.items(): + previous[sig] = signal.getsignal(sig) + signal.signal(sig, self._handler_for(disposition)) + return previous + + def _restore_handlers(self, previous: dict[int, object]) -> None: + for sig, handler in previous.items(): + signal.signal(sig, handler) + + def _handler_for(self, disposition: Disposition): + if disposition == "teardown": + + def teardown_handler(signum, frame): # noqa: ARG001 + raise KeyboardInterrupt + + return teardown_handler + + def detach_handler(signum, frame): # noqa: ARG001 + raise SystemExit(0) + + return detach_handler diff --git a/areal/experimental/inference_service/__init__.py b/areal/v2/inference_service/__init__.py similarity index 60% rename from areal/experimental/inference_service/__init__.py rename to areal/v2/inference_service/__init__.py index edbc897123..1517b17181 100644 --- a/areal/experimental/inference_service/__init__.py +++ b/areal/v2/inference_service/__init__.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 -from areal.experimental.inference_service.controller.workflow import ( +from areal.v2.inference_service.controller.workflow import ( InferenceServiceWorkflow, ) diff --git a/areal/experimental/inference_service/backend.py b/areal/v2/inference_service/backend.py similarity index 100% rename from areal/experimental/inference_service/backend.py rename to areal/v2/inference_service/backend.py diff --git a/areal/experimental/inference_service/controller/__init__.py b/areal/v2/inference_service/controller/__init__.py similarity index 71% rename from areal/experimental/inference_service/controller/__init__.py rename to areal/v2/inference_service/controller/__init__.py index 9ac57d0d68..63c6127d57 100644 --- a/areal/experimental/inference_service/controller/__init__.py +++ b/areal/v2/inference_service/controller/__init__.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 from areal.api.cli_args import InferenceEngineConfig -from areal.experimental.inference_service.controller.controller import ( +from areal.v2.inference_service.controller.controller import ( RolloutControllerV2, ) diff --git a/areal/experimental/inference_service/controller/controller.py b/areal/v2/inference_service/controller/controller.py similarity index 99% rename from areal/experimental/inference_service/controller/controller.py rename to areal/v2/inference_service/controller/controller.py index 6f45861c3b..40e3d13831 100644 --- a/areal/experimental/inference_service/controller/controller.py +++ b/areal/v2/inference_service/controller/controller.py @@ -338,7 +338,7 @@ async def _async_initialize( port_count=2, gpu=0, mem=8, - cmd="python -m areal.experimental.inference_service.guard", + cmd="python -m areal.v2.inference_service.guard", ) total_workers = dp_size else: @@ -357,7 +357,7 @@ async def _async_initialize( if inf_spec.gpu > 0: inf_spec.gpu = gpus_per_worker - inf_spec.cmd = "python -m areal.experimental.inference_service.guard" + inf_spec.cmd = "python -m areal.v2.inference_service.guard" inf_role = f"{self._worker_role}{self._INF_SUFFIX}" inf_job = Job( @@ -392,7 +392,7 @@ async def _async_initialize( router_cmd = [ sys.executable, "-m", - "areal.experimental.inference_service.router", + "areal.v2.inference_service.router", "--admin-api-key", admin_api_key, "--routing-strategy", @@ -451,7 +451,7 @@ async def _async_initialize( data_proxy_base_cmd = [ sys.executable, "-m", - "areal.experimental.inference_service.data_proxy", + "areal.v2.inference_service.data_proxy", "--tokenizer-path", cfg.tokenizer_path, "--admin-api-key", @@ -507,7 +507,7 @@ async def _fork_data_proxy(group_idx: int) -> tuple[str, int, str]: gw_cmd = [ sys.executable, "-m", - "areal.experimental.inference_service.gateway", + "areal.v2.inference_service.gateway", "--admin-api-key", admin_api_key, "--router-addr", @@ -1523,7 +1523,7 @@ def _wrap_agent(self, agent: Any, group_size: int = 1): group_size : int Number of parallel trajectories per episode. """ - from areal.experimental.inference_service.controller.workflow import ( + from areal.v2.inference_service.controller.workflow import ( InferenceServiceWorkflow, ) @@ -1593,7 +1593,7 @@ def _resolve_workflow( f"Got group_size={group_size}." ) - from areal.experimental.inference_service.controller.workflow import ( + from areal.v2.inference_service.controller.workflow import ( InferenceServiceWorkflow, ) diff --git a/areal/experimental/inference_service/controller/workflow.py b/areal/v2/inference_service/controller/workflow.py similarity index 99% rename from areal/experimental/inference_service/controller/workflow.py rename to areal/v2/inference_service/controller/workflow.py index 98be7bb1b8..dd51256ba7 100644 --- a/areal/experimental/inference_service/controller/workflow.py +++ b/areal/v2/inference_service/controller/workflow.py @@ -18,10 +18,10 @@ if TYPE_CHECKING: from areal.api.engine_api import InferenceEngine - from areal.experimental.inference_service.controller.controller import ( + from areal.experimental.openai.types import InteractionWithTokenLogpReward + from areal.v2.inference_service.controller.controller import ( RolloutControllerV2, ) - from areal.experimental.openai.types import InteractionWithTokenLogpReward logger = logging.getLogger("InferenceServiceWorkflow") diff --git a/areal/experimental/training_service/controller/__init__.py b/areal/v2/inference_service/data_proxy/__init__.py similarity index 100% rename from areal/experimental/training_service/controller/__init__.py rename to areal/v2/inference_service/data_proxy/__init__.py diff --git a/areal/experimental/inference_service/data_proxy/__main__.py b/areal/v2/inference_service/data_proxy/__main__.py similarity index 92% rename from areal/experimental/inference_service/data_proxy/__main__.py rename to areal/v2/inference_service/data_proxy/__main__.py index 010c393e36..e0458bebf9 100644 --- a/areal/experimental/inference_service/data_proxy/__main__.py +++ b/areal/v2/inference_service/data_proxy/__main__.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 -"""CLI entrypoint: python -m areal.experimental.inference_service.data_proxy""" +"""CLI entrypoint: python -m areal.v2.inference_service.data_proxy""" from __future__ import annotations @@ -8,14 +8,14 @@ import uvicorn -from areal.experimental.inference_service.data_proxy.app import create_app -from areal.experimental.inference_service.data_proxy.config import DataProxyConfig from areal.infra.utils.http import ( get_default_uvicorn_kwargs, validate_admin_api_key, ) from areal.utils.logging import suppress_http_loggers from areal.utils.network import format_hostport +from areal.v2.inference_service.data_proxy.app import create_app +from areal.v2.inference_service.data_proxy.config import DataProxyConfig def main(): diff --git a/areal/experimental/inference_service/data_proxy/app.py b/areal/v2/inference_service/data_proxy/app.py similarity index 98% rename from areal/experimental/inference_service/data_proxy/app.py rename to areal/v2/inference_service/data_proxy/app.py index fd2794983c..5f703d7e86 100644 --- a/areal/experimental/inference_service/data_proxy/app.py +++ b/areal/v2/inference_service/data_proxy/app.py @@ -17,25 +17,6 @@ from flask import Flask from pydantic import BaseModel -from areal.experimental.inference_service.data_proxy.config import DataProxyConfig -from areal.experimental.inference_service.data_proxy.pause import PauseState -from areal.experimental.inference_service.data_proxy.session import ( - ExportTrajectoriesRequest, - ExportTrajectoriesResponse, - ReadyNotification, - SessionCredentials, - SessionData, - SessionStore, - SetRewardRequest, - StartSessionRequest, - StartSessionResponse, -) -from areal.experimental.inference_service.data_proxy.tokenizer_proxy import ( - TokenizerProxy, -) -from areal.experimental.inference_service.inf_bridge import InfBridge -from areal.experimental.inference_service.sglang.bridge import SGLangBridgeBackend -from areal.experimental.inference_service.vllm.bridge import VLLMBridgeBackend from areal.experimental.openai.client import ArealOpenAI from areal.experimental.openai.types import ( InteractionWithTokenLogpReward, @@ -49,6 +30,25 @@ from areal.infra.utils.http import create_httpx_client from areal.utils import logging from areal.utils.data import concat_padded_tensors +from areal.v2.inference_service.data_proxy.config import DataProxyConfig +from areal.v2.inference_service.data_proxy.pause import PauseState +from areal.v2.inference_service.data_proxy.session import ( + ExportTrajectoriesRequest, + ExportTrajectoriesResponse, + ReadyNotification, + SessionCredentials, + SessionData, + SessionStore, + SetRewardRequest, + StartSessionRequest, + StartSessionResponse, +) +from areal.v2.inference_service.data_proxy.tokenizer_proxy import ( + TokenizerProxy, +) +from areal.v2.inference_service.inf_bridge import InfBridge +from areal.v2.inference_service.sglang.bridge import SGLangBridgeBackend +from areal.v2.inference_service.vllm.bridge import VLLMBridgeBackend logger = logging.getLogger("InferenceDataProxy") diff --git a/areal/experimental/inference_service/data_proxy/config.py b/areal/v2/inference_service/data_proxy/config.py similarity index 100% rename from areal/experimental/inference_service/data_proxy/config.py rename to areal/v2/inference_service/data_proxy/config.py diff --git a/areal/experimental/inference_service/data_proxy/pause.py b/areal/v2/inference_service/data_proxy/pause.py similarity index 100% rename from areal/experimental/inference_service/data_proxy/pause.py rename to areal/v2/inference_service/data_proxy/pause.py diff --git a/areal/experimental/inference_service/data_proxy/session.py b/areal/v2/inference_service/data_proxy/session.py similarity index 100% rename from areal/experimental/inference_service/data_proxy/session.py rename to areal/v2/inference_service/data_proxy/session.py diff --git a/areal/experimental/inference_service/data_proxy/tokenizer_proxy.py b/areal/v2/inference_service/data_proxy/tokenizer_proxy.py similarity index 100% rename from areal/experimental/inference_service/data_proxy/tokenizer_proxy.py rename to areal/v2/inference_service/data_proxy/tokenizer_proxy.py diff --git a/areal/experimental/training_service/data_proxy/__init__.py b/areal/v2/inference_service/gateway/__init__.py similarity index 100% rename from areal/experimental/training_service/data_proxy/__init__.py rename to areal/v2/inference_service/gateway/__init__.py diff --git a/areal/experimental/inference_service/gateway/__main__.py b/areal/v2/inference_service/gateway/__main__.py similarity index 88% rename from areal/experimental/inference_service/gateway/__main__.py rename to areal/v2/inference_service/gateway/__main__.py index 1d6845d91d..c5b278d8a6 100644 --- a/areal/experimental/inference_service/gateway/__main__.py +++ b/areal/v2/inference_service/gateway/__main__.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 -"""CLI entrypoint for the inference gateway: ``python -m areal.experimental.inference_service.gateway``.""" +"""CLI entrypoint for the inference gateway: ``python -m areal.v2.inference_service.gateway``.""" from __future__ import annotations @@ -41,13 +41,13 @@ def main(): ) args, _ = parser.parse_known_args() - from areal.experimental.inference_service.gateway.app import create_app - from areal.experimental.inference_service.gateway.config import GatewayConfig from areal.infra.utils.http import ( get_default_uvicorn_kwargs, validate_admin_api_key, ) from areal.utils.logging import suppress_http_loggers + from areal.v2.inference_service.gateway.app import create_app + from areal.v2.inference_service.gateway.config import GatewayConfig validate_admin_api_key(args.host, args.admin_api_key) diff --git a/areal/experimental/inference_service/gateway/app.py b/areal/v2/inference_service/gateway/app.py similarity index 99% rename from areal/experimental/inference_service/gateway/app.py rename to areal/v2/inference_service/gateway/app.py index 902db4caad..a908780f2f 100644 --- a/areal/experimental/inference_service/gateway/app.py +++ b/areal/v2/inference_service/gateway/app.py @@ -18,12 +18,14 @@ from fastapi.responses import JSONResponse, Response, StreamingResponse from pydantic import BaseModel -from areal.experimental.inference_service.gateway.auth import ( +from areal.infra.utils.http import create_httpx_client +from areal.utils import logging +from areal.v2.inference_service.gateway.auth import ( extract_bearer_token, require_admin_key, ) -from areal.experimental.inference_service.gateway.config import GatewayConfig -from areal.experimental.inference_service.gateway.streaming import ( +from areal.v2.inference_service.gateway.config import GatewayConfig +from areal.v2.inference_service.gateway.streaming import ( RouterKeyRejectedError, RouterUnreachableError, _forwarding_headers, @@ -38,8 +40,6 @@ resolve_worker_addr, revoke_session_in_router, ) -from areal.infra.utils.http import create_httpx_client -from areal.utils import logging logger = logging.getLogger("InferenceGateway") diff --git a/areal/experimental/inference_service/gateway/auth.py b/areal/v2/inference_service/gateway/auth.py similarity index 100% rename from areal/experimental/inference_service/gateway/auth.py rename to areal/v2/inference_service/gateway/auth.py diff --git a/areal/experimental/inference_service/gateway/config.py b/areal/v2/inference_service/gateway/config.py similarity index 100% rename from areal/experimental/inference_service/gateway/config.py rename to areal/v2/inference_service/gateway/config.py diff --git a/areal/experimental/inference_service/gateway/streaming.py b/areal/v2/inference_service/gateway/streaming.py similarity index 100% rename from areal/experimental/inference_service/gateway/streaming.py rename to areal/v2/inference_service/gateway/streaming.py diff --git a/areal/experimental/training_service/gateway/__init__.py b/areal/v2/inference_service/guard/__init__.py similarity index 100% rename from areal/experimental/training_service/gateway/__init__.py rename to areal/v2/inference_service/guard/__init__.py diff --git a/areal/experimental/inference_service/guard/__main__.py b/areal/v2/inference_service/guard/__main__.py similarity index 79% rename from areal/experimental/inference_service/guard/__main__.py rename to areal/v2/inference_service/guard/__main__.py index 8e3a806b28..7710f46611 100644 --- a/areal/experimental/inference_service/guard/__main__.py +++ b/areal/v2/inference_service/guard/__main__.py @@ -1,18 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 -"""CLI entrypoint: ``python -m areal.experimental.inference_service.guard``""" +"""CLI entrypoint: ``python -m areal.v2.inference_service.guard``""" from __future__ import annotations -from areal.experimental.inference_service.guard.app import ( - _state, - app, -) from areal.infra.rpc.guard.app import ( configure_state_from_args, make_base_parser, run_server, ) +from areal.v2.inference_service.guard.app import ( + _state, + app, +) def main(): diff --git a/areal/experimental/inference_service/guard/app.py b/areal/v2/inference_service/guard/app.py similarity index 100% rename from areal/experimental/inference_service/guard/app.py rename to areal/v2/inference_service/guard/app.py diff --git a/areal/experimental/inference_service/inf_bridge.py b/areal/v2/inference_service/inf_bridge.py similarity index 98% rename from areal/experimental/inference_service/inf_bridge.py rename to areal/v2/inference_service/inf_bridge.py index f7532941f2..94ce45b118 100644 --- a/areal/experimental/inference_service/inf_bridge.py +++ b/areal/v2/inference_service/inf_bridge.py @@ -17,12 +17,12 @@ import numpy as np from areal.api.io_struct import HttpRequest -from areal.experimental.inference_service.backend import InfBridgeBackend from areal.utils import logging +from areal.v2.inference_service.backend import InfBridgeBackend if TYPE_CHECKING: from areal.api.io_struct import ModelRequest, ModelResponse - from areal.experimental.inference_service.data_proxy.pause import PauseState + from areal.v2.inference_service.data_proxy.pause import PauseState _StopReason = Literal["length", "stop", "tool_calls", "abort"] diff --git a/areal/experimental/training_service/router/__init__.py b/areal/v2/inference_service/router/__init__.py similarity index 100% rename from areal/experimental/training_service/router/__init__.py rename to areal/v2/inference_service/router/__init__.py diff --git a/areal/experimental/inference_service/router/__main__.py b/areal/v2/inference_service/router/__main__.py similarity index 88% rename from areal/experimental/inference_service/router/__main__.py rename to areal/v2/inference_service/router/__main__.py index d2502bb146..5557abd812 100644 --- a/areal/experimental/inference_service/router/__main__.py +++ b/areal/v2/inference_service/router/__main__.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 -"""CLI entrypoint for the router: ``python -m areal.experimental.inference_service.router``.""" +"""CLI entrypoint for the router: ``python -m areal.v2.inference_service.router``.""" from __future__ import annotations @@ -42,13 +42,13 @@ def main(): ) args, _ = parser.parse_known_args() - from areal.experimental.inference_service.router.app import create_app - from areal.experimental.inference_service.router.config import RouterConfig from areal.infra.utils.http import ( get_default_uvicorn_kwargs, validate_admin_api_key, ) from areal.utils.logging import suppress_http_loggers + from areal.v2.inference_service.router.app import create_app + from areal.v2.inference_service.router.config import RouterConfig validate_admin_api_key(args.host, args.admin_api_key) diff --git a/areal/experimental/inference_service/router/app.py b/areal/v2/inference_service/router/app.py similarity index 98% rename from areal/experimental/inference_service/router/app.py rename to areal/v2/inference_service/router/app.py index e8f5cb8b0a..6274327872 100644 --- a/areal/experimental/inference_service/router/app.py +++ b/areal/v2/inference_service/router/app.py @@ -7,7 +7,7 @@ It never proxies traffic — it only answers routing queries. Endpoint names are aligned with -``areal.experimental.agent_service.router.app``: +``areal.v2.agent_service.router.app``: ``/register``, ``/unregister``, ``/route``, ``/remove_session``. """ @@ -21,15 +21,15 @@ from fastapi import FastAPI, HTTPException, Request from pydantic import BaseModel -from areal.experimental.inference_service.router.config import RouterConfig -from areal.experimental.inference_service.router.state import ( +from areal.utils import logging +from areal.v2.inference_service.router.config import RouterConfig +from areal.v2.inference_service.router.state import ( GroupRegistry, ModelRegistry, SessionRegistry, WorkerRegistry, ) -from areal.experimental.inference_service.router.strategies import get_strategy -from areal.utils import logging +from areal.v2.inference_service.router.strategies import get_strategy logger = logging.getLogger("InferenceRouter") diff --git a/areal/experimental/inference_service/router/config.py b/areal/v2/inference_service/router/config.py similarity index 100% rename from areal/experimental/inference_service/router/config.py rename to areal/v2/inference_service/router/config.py diff --git a/areal/experimental/inference_service/router/state.py b/areal/v2/inference_service/router/state.py similarity index 100% rename from areal/experimental/inference_service/router/state.py rename to areal/v2/inference_service/router/state.py diff --git a/areal/experimental/inference_service/router/strategies.py b/areal/v2/inference_service/router/strategies.py similarity index 94% rename from areal/experimental/inference_service/router/strategies.py rename to areal/v2/inference_service/router/strategies.py index d76e2e87b9..3c7d7f8f3b 100644 --- a/areal/experimental/inference_service/router/strategies.py +++ b/areal/v2/inference_service/router/strategies.py @@ -6,7 +6,7 @@ from typing import Protocol -from areal.experimental.inference_service.router.state import WorkerInfo +from areal.v2.inference_service.router.state import WorkerInfo class RoutingStrategy(Protocol): diff --git a/areal/experimental/training_service/worker/__init__.py b/areal/v2/inference_service/sglang/__init__.py similarity index 100% rename from areal/experimental/training_service/worker/__init__.py rename to areal/v2/inference_service/sglang/__init__.py diff --git a/areal/experimental/inference_service/sglang/awex.py b/areal/v2/inference_service/sglang/awex.py similarity index 98% rename from areal/experimental/inference_service/sglang/awex.py rename to areal/v2/inference_service/sglang/awex.py index 6bd3f3c015..ceef2778f8 100644 --- a/areal/experimental/inference_service/sglang/awex.py +++ b/areal/v2/inference_service/sglang/awex.py @@ -9,7 +9,7 @@ from areal.utils import logging if TYPE_CHECKING: - from areal.experimental.inference_service.sglang.rpc_proxy import RpcProxy + from areal.v2.inference_service.sglang.rpc_proxy import RpcProxy logger = logging.getLogger("AwexInferenceEndpoints") diff --git a/areal/experimental/inference_service/sglang/bridge.py b/areal/v2/inference_service/sglang/bridge.py similarity index 100% rename from areal/experimental/inference_service/sglang/bridge.py rename to areal/v2/inference_service/sglang/bridge.py diff --git a/areal/experimental/inference_service/sglang/launch_server.py b/areal/v2/inference_service/sglang/launch_server.py similarity index 92% rename from areal/experimental/inference_service/sglang/launch_server.py rename to areal/v2/inference_service/sglang/launch_server.py index 1894398674..0406b2ba64 100644 --- a/areal/experimental/inference_service/sglang/launch_server.py +++ b/areal/v2/inference_service/sglang/launch_server.py @@ -24,11 +24,11 @@ def areal_launch_server(server_args) -> None: from sglang.srt.managers.detokenizer_manager import run_detokenizer_process # ---- BEGIN AREAL ---- - from areal.experimental.inference_service.sglang.awex import ( + from areal.v2.inference_service.sglang.awex import ( register_awex_endpoints, ) - from areal.experimental.inference_service.sglang.rpc_proxy import RpcProxy - from areal.experimental.inference_service.sglang.scheduler import ( + from areal.v2.inference_service.sglang.rpc_proxy import RpcProxy + from areal.v2.inference_service.sglang.scheduler import ( areal_run_scheduler_process, create_result_ipc, ) diff --git a/areal/experimental/inference_service/sglang/pp_bridge.py b/areal/v2/inference_service/sglang/pp_bridge.py similarity index 100% rename from areal/experimental/inference_service/sglang/pp_bridge.py rename to areal/v2/inference_service/sglang/pp_bridge.py diff --git a/areal/experimental/inference_service/sglang/rpc_proxy.py b/areal/v2/inference_service/sglang/rpc_proxy.py similarity index 100% rename from areal/experimental/inference_service/sglang/rpc_proxy.py rename to areal/v2/inference_service/sglang/rpc_proxy.py diff --git a/areal/experimental/inference_service/sglang/scheduler.py b/areal/v2/inference_service/sglang/scheduler.py similarity index 98% rename from areal/experimental/inference_service/sglang/scheduler.py rename to areal/v2/inference_service/sglang/scheduler.py index 34149a6fc9..57ce2db7c7 100644 --- a/areal/experimental/inference_service/sglang/scheduler.py +++ b/areal/v2/inference_service/sglang/scheduler.py @@ -72,7 +72,7 @@ def bind(self) -> None: def _require_adapter(self) -> Any: if self._adapter is None: - from areal.experimental.weight_update.awex.sglang_adapter import ( + from areal.v2.weight_update.awex.sglang_adapter import ( AwexSGLangAdapter, ) @@ -186,7 +186,7 @@ def areal_run_scheduler_process( ) from sglang.utils import get_exception_traceback - from areal.experimental.inference_service.sglang.pp_bridge import ( + from areal.v2.inference_service.sglang.pp_bridge import ( PPSchedulerBridge, ) diff --git a/areal/experimental/weight_update/gateway/__init__.py b/areal/v2/inference_service/vllm/__init__.py similarity index 100% rename from areal/experimental/weight_update/gateway/__init__.py rename to areal/v2/inference_service/vllm/__init__.py diff --git a/areal/experimental/inference_service/vllm/bridge.py b/areal/v2/inference_service/vllm/bridge.py similarity index 100% rename from areal/experimental/inference_service/vllm/bridge.py rename to areal/v2/inference_service/vllm/bridge.py diff --git a/areal/experimental/training_service/__init__.py b/areal/v2/training_service/__init__.py similarity index 100% rename from areal/experimental/training_service/__init__.py rename to areal/v2/training_service/__init__.py diff --git a/tests/experimental/weight_update/__init__.py b/areal/v2/training_service/controller/__init__.py similarity index 100% rename from tests/experimental/weight_update/__init__.py rename to areal/v2/training_service/controller/__init__.py diff --git a/areal/experimental/training_service/controller/controller.py b/areal/v2/training_service/controller/controller.py similarity index 94% rename from areal/experimental/training_service/controller/controller.py rename to areal/v2/training_service/controller/controller.py index 462e73d584..16f136095e 100644 --- a/areal/experimental/training_service/controller/controller.py +++ b/areal/v2/training_service/controller/controller.py @@ -17,7 +17,7 @@ from areal.infra.utils.concurrent import get_executor, run_async_task from areal.infra.utils.http import create_httpx_client from areal.utils import logging -from areal.utils.network import format_hostport +from areal.utils.network import format_hostport, gethostip if TYPE_CHECKING: from areal.api import ParallelStrategy, TrainEngine @@ -178,7 +178,7 @@ async def _async_initialize( ) guard_spec = SchedulingSpec(**asdict(cfg.scheduling_spec[0])) - guard_spec.cmd = "python -m areal.experimental.training_service.guard" + guard_spec.cmd = "python -m areal.v2.training_service.guard" guard_role = f"{role}{self._GUARD_SUFFIX}" guard_job = Job( @@ -243,7 +243,7 @@ async def _fork_worker(rank: int) -> str: worker_cmd = [ sys.executable, "-m", - "areal.experimental.training_service.worker", + "areal.v2.training_service.worker", "--admin-api-key", cfg.admin_api_key, "--log-level", @@ -327,7 +327,7 @@ async def _fork_worker(rank: int) -> str: router_cmd = [ sys.executable, "-m", - "areal.experimental.training_service.router", + "areal.v2.training_service.router", "--admin-api-key", cfg.admin_api_key, "--log-level", @@ -351,7 +351,7 @@ async def _fork_worker(rank: int) -> str: data_proxy_cmd = [ sys.executable, "-m", - "areal.experimental.training_service.data_proxy", + "areal.v2.training_service.data_proxy", "--worker-addrs", ",".join(self._worker_addrs), "--admin-api-key", @@ -377,7 +377,7 @@ async def _fork_worker(rank: int) -> str: gw_cmd = [ sys.executable, "-m", - "areal.experimental.training_service.gateway", + "areal.v2.training_service.gateway", "--admin-api-key", cfg.admin_api_key, "--router-addr", @@ -936,13 +936,13 @@ def connect_engine(self, rollout: Any, meta: Any) -> None: self._ensure_initialized() import requests - from areal.experimental.inference_service.controller.controller import ( + from areal.v2.inference_service.controller.controller import ( RolloutControllerV2, ) - from areal.experimental.weight_update.controller.config import ( + from areal.v2.weight_update.controller.config import ( WeightUpdateControllerConfig, ) - from areal.experimental.weight_update.controller.controller import ( + from areal.v2.weight_update.controller.controller import ( WeightUpdateController, ) @@ -955,13 +955,18 @@ def connect_engine(self, rollout: Any, meta: Any) -> None: self.rollout = rollout - if meta.type != "awex": + if meta.type not in ("awex", "disk"): raise ValueError( - f"GatewayTrainController only supports 'awex' weight updates, got '{meta.type}'" + f"GatewayTrainController supports 'awex' or 'disk' weight " + f"updates, got '{meta.type}'" ) ctrl = WeightUpdateController( WeightUpdateControllerConfig( + # Bind gateway to this node's outbound IP so cross-host + # train/inf workers can reach the kv_store_url the gateway + # publishes. Default "127.0.0.1" only works single-machine. + host=gethostip(), admin_api_key=self.config.admin_api_key, log_level=self.config.log_level, ) @@ -969,29 +974,39 @@ def connect_engine(self, rollout: Any, meta: Any) -> None: ctrl.initialize() inference_urls: list[str] = rollout.inference_worker_urls - - # NCCL rendezvous master must live on the rank-0 process's node. - # awex assigns rank 0 to inference[0], so allocate on the inference - # rank-0 guard rather than a train guard. - inf_guard_addrs = rollout.inference_guard_addrs - resp = requests.post( - f"{inf_guard_addrs[0]}/alloc_ports", - json={"count": 1}, - timeout=30, - ) - resp.raise_for_status() - port_data = resp.json() - nccl_master_addr = port_data["host"] - nccl_master_port = port_data["ports"][0] - pair_name = f"{self._role}-rollout" - ctrl.connect( - pair_name=pair_name, - train_worker_urls=self._worker_addrs, - inference_worker_urls=inference_urls, - nccl_master_addr=nccl_master_addr, - nccl_master_port=nccl_master_port, - ) + + if meta.type == "awex": + # NCCL rendezvous master must live on the rank-0 process's node. + # awex assigns rank 0 to inference[0], so allocate on the inference + # rank-0 guard rather than a train guard. + inf_guard_addrs = rollout.inference_guard_addrs + resp = requests.post( + f"{inf_guard_addrs[0]}/alloc_ports", + json={"count": 1}, + timeout=30, + ) + resp.raise_for_status() + port_data = resp.json() + ctrl.connect( + pair_name=pair_name, + train_worker_urls=self._worker_addrs, + inference_worker_urls=inference_urls, + mode="awex", + nccl_master_addr=port_data["host"], + nccl_master_port=port_data["ports"][0], + ) + else: # disk + ctrl.connect( + pair_name=pair_name, + train_worker_urls=self._worker_addrs, + inference_worker_urls=inference_urls, + mode="disk", + save_path=meta.path or "", + use_lora=meta.use_lora, + lora_name=meta.lora_name, + lora_keep_versions=meta.lora_keep_versions, + ) self._weight_update_ctrl = ctrl logger.info( "WeightUpdateController connected (pair=%s, train=%d, inf=%d)", diff --git a/tests/experimental/weight_update/torchrun/__init__.py b/areal/v2/training_service/data_proxy/__init__.py similarity index 100% rename from tests/experimental/weight_update/torchrun/__init__.py rename to areal/v2/training_service/data_proxy/__init__.py diff --git a/areal/experimental/training_service/data_proxy/__main__.py b/areal/v2/training_service/data_proxy/__main__.py similarity index 91% rename from areal/experimental/training_service/data_proxy/__main__.py rename to areal/v2/training_service/data_proxy/__main__.py index 70e18c2b08..82ff2cb515 100644 --- a/areal/experimental/training_service/data_proxy/__main__.py +++ b/areal/v2/training_service/data_proxy/__main__.py @@ -6,13 +6,13 @@ import uvicorn -from areal.experimental.training_service.data_proxy.app import create_app -from areal.experimental.training_service.data_proxy.config import TrainDataProxyConfig from areal.infra.utils.http import ( get_default_uvicorn_kwargs, validate_admin_api_key, ) from areal.utils.logging import suppress_http_loggers +from areal.v2.training_service.data_proxy.app import create_app +from areal.v2.training_service.data_proxy.config import TrainDataProxyConfig def main(): diff --git a/areal/experimental/training_service/data_proxy/app.py b/areal/v2/training_service/data_proxy/app.py similarity index 78% rename from areal/experimental/training_service/data_proxy/app.py rename to areal/v2/training_service/data_proxy/app.py index ac888918fb..55488f6bca 100644 --- a/areal/experimental/training_service/data_proxy/app.py +++ b/areal/v2/training_service/data_proxy/app.py @@ -7,11 +7,11 @@ from fastapi import FastAPI from fastapi.responses import Response as RawResponse -from areal.experimental.training_service.data_proxy.config import TrainDataProxyConfig -from areal.experimental.training_service.data_proxy.dispatcher import Dispatcher -from areal.experimental.training_service.data_proxy.engine import register_engine_routes -from areal.experimental.training_service.data_proxy.topology import discover_topology from areal.utils import logging +from areal.v2.training_service.data_proxy.config import TrainDataProxyConfig +from areal.v2.training_service.data_proxy.dispatcher import Dispatcher +from areal.v2.training_service.data_proxy.engine import register_engine_routes +from areal.v2.training_service.data_proxy.topology import discover_topology logger = logging.getLogger("TrainDataProxy") diff --git a/areal/experimental/training_service/data_proxy/config.py b/areal/v2/training_service/data_proxy/config.py similarity index 100% rename from areal/experimental/training_service/data_proxy/config.py rename to areal/v2/training_service/data_proxy/config.py diff --git a/areal/experimental/training_service/data_proxy/dispatcher.py b/areal/v2/training_service/data_proxy/dispatcher.py similarity index 99% rename from areal/experimental/training_service/data_proxy/dispatcher.py rename to areal/v2/training_service/data_proxy/dispatcher.py index 87df614c42..7f829a74b6 100644 --- a/areal/experimental/training_service/data_proxy/dispatcher.py +++ b/areal/v2/training_service/data_proxy/dispatcher.py @@ -34,7 +34,6 @@ import aiohttp import orjson -from areal.experimental.training_service.data_proxy.topology import WorkerTopology from areal.infra.controller.train_controller import ( _dispatch_tensors, _is_tensor_like, @@ -43,6 +42,7 @@ ) from areal.infra.rpc.serialization import deserialize_value, serialize_value from areal.utils import logging +from areal.v2.training_service.data_proxy.topology import WorkerTopology logger = logging.getLogger("TrainDataProxy") diff --git a/areal/experimental/training_service/data_proxy/engine.py b/areal/v2/training_service/data_proxy/engine.py similarity index 100% rename from areal/experimental/training_service/data_proxy/engine.py rename to areal/v2/training_service/data_proxy/engine.py diff --git a/areal/experimental/training_service/data_proxy/topology.py b/areal/v2/training_service/data_proxy/topology.py similarity index 100% rename from areal/experimental/training_service/data_proxy/topology.py rename to areal/v2/training_service/data_proxy/topology.py diff --git a/areal/v2/training_service/gateway/__init__.py b/areal/v2/training_service/gateway/__init__.py new file mode 100644 index 0000000000..9881313609 --- /dev/null +++ b/areal/v2/training_service/gateway/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/areal/experimental/training_service/gateway/__main__.py b/areal/v2/training_service/gateway/__main__.py similarity index 92% rename from areal/experimental/training_service/gateway/__main__.py rename to areal/v2/training_service/gateway/__main__.py index a4cdfdff16..8940417e74 100644 --- a/areal/experimental/training_service/gateway/__main__.py +++ b/areal/v2/training_service/gateway/__main__.py @@ -39,13 +39,13 @@ def main(): ) args, _ = parser.parse_known_args() - from areal.experimental.training_service.gateway.app import create_app - from areal.experimental.training_service.gateway.config import GatewayConfig from areal.infra.utils.http import ( get_default_uvicorn_kwargs, validate_admin_api_key, ) from areal.utils.logging import suppress_http_loggers + from areal.v2.training_service.gateway.app import create_app + from areal.v2.training_service.gateway.config import GatewayConfig validate_admin_api_key(args.host, args.admin_api_key) diff --git a/areal/experimental/training_service/gateway/app.py b/areal/v2/training_service/gateway/app.py similarity index 92% rename from areal/experimental/training_service/gateway/app.py rename to areal/v2/training_service/gateway/app.py index f539a28b10..018273f1d8 100644 --- a/areal/experimental/training_service/gateway/app.py +++ b/areal/v2/training_service/gateway/app.py @@ -8,11 +8,11 @@ from fastapi import FastAPI, Request from fastapi.responses import JSONResponse, Response -from areal.experimental.training_service.gateway import streaming -from areal.experimental.training_service.gateway.auth import extract_bearer_token -from areal.experimental.training_service.gateway.config import GatewayConfig -from areal.experimental.training_service.gateway.engine import register_engine_routes from areal.utils import logging +from areal.v2.training_service.gateway import streaming +from areal.v2.training_service.gateway.auth import extract_bearer_token +from areal.v2.training_service.gateway.config import GatewayConfig +from areal.v2.training_service.gateway.engine import register_engine_routes logger = logging.getLogger("TrainGateway") diff --git a/areal/experimental/training_service/gateway/auth.py b/areal/v2/training_service/gateway/auth.py similarity index 100% rename from areal/experimental/training_service/gateway/auth.py rename to areal/v2/training_service/gateway/auth.py diff --git a/areal/experimental/training_service/gateway/config.py b/areal/v2/training_service/gateway/config.py similarity index 100% rename from areal/experimental/training_service/gateway/config.py rename to areal/v2/training_service/gateway/config.py diff --git a/areal/experimental/training_service/gateway/engine.py b/areal/v2/training_service/gateway/engine.py similarity index 98% rename from areal/experimental/training_service/gateway/engine.py rename to areal/v2/training_service/gateway/engine.py index 736b5e4968..f14c195069 100644 --- a/areal/experimental/training_service/gateway/engine.py +++ b/areal/v2/training_service/gateway/engine.py @@ -7,7 +7,7 @@ from fastapi import FastAPI, Request -from areal.experimental.training_service.gateway.config import GatewayConfig +from areal.v2.training_service.gateway.config import GatewayConfig def register_engine_routes( diff --git a/areal/experimental/training_service/gateway/streaming.py b/areal/v2/training_service/gateway/streaming.py similarity index 100% rename from areal/experimental/training_service/gateway/streaming.py rename to areal/v2/training_service/gateway/streaming.py diff --git a/areal/experimental/training_service/guard/__init__.py b/areal/v2/training_service/guard/__init__.py similarity index 100% rename from areal/experimental/training_service/guard/__init__.py rename to areal/v2/training_service/guard/__init__.py diff --git a/areal/experimental/training_service/guard/__main__.py b/areal/v2/training_service/guard/__main__.py similarity index 77% rename from areal/experimental/training_service/guard/__main__.py rename to areal/v2/training_service/guard/__main__.py index 3028624721..02d648f781 100644 --- a/areal/experimental/training_service/guard/__main__.py +++ b/areal/v2/training_service/guard/__main__.py @@ -1,15 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 -"""CLI entrypoint: ``python -m areal.experimental.training_service.guard``""" +"""CLI entrypoint: ``python -m areal.v2.training_service.guard``""" from __future__ import annotations -from areal.experimental.training_service.guard.app import _state, app from areal.infra.rpc.guard.app import ( configure_state_from_args, make_base_parser, run_server, ) +from areal.v2.training_service.guard.app import _state, app def main() -> None: diff --git a/areal/experimental/training_service/guard/app.py b/areal/v2/training_service/guard/app.py similarity index 100% rename from areal/experimental/training_service/guard/app.py rename to areal/v2/training_service/guard/app.py diff --git a/areal/v2/training_service/router/__init__.py b/areal/v2/training_service/router/__init__.py new file mode 100644 index 0000000000..9881313609 --- /dev/null +++ b/areal/v2/training_service/router/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/areal/experimental/training_service/router/__main__.py b/areal/v2/training_service/router/__main__.py similarity index 92% rename from areal/experimental/training_service/router/__main__.py rename to areal/v2/training_service/router/__main__.py index 3eecfda2e9..bcbbb16372 100644 --- a/areal/experimental/training_service/router/__main__.py +++ b/areal/v2/training_service/router/__main__.py @@ -35,13 +35,13 @@ def main(): ) args, _ = parser.parse_known_args() - from areal.experimental.training_service.router.app import create_app - from areal.experimental.training_service.router.config import RouterConfig from areal.infra.utils.http import ( get_default_uvicorn_kwargs, validate_admin_api_key, ) from areal.utils.logging import suppress_http_loggers + from areal.v2.training_service.router.app import create_app + from areal.v2.training_service.router.config import RouterConfig validate_admin_api_key(args.host, args.admin_api_key) diff --git a/areal/experimental/training_service/router/app.py b/areal/v2/training_service/router/app.py similarity index 97% rename from areal/experimental/training_service/router/app.py rename to areal/v2/training_service/router/app.py index ad8b1aa257..1cca634595 100644 --- a/areal/experimental/training_service/router/app.py +++ b/areal/v2/training_service/router/app.py @@ -8,9 +8,9 @@ from contextlib import asynccontextmanager from typing import Any -from areal.experimental.training_service.router.config import RouterConfig -from areal.experimental.training_service.router.state import ModelRegistry from areal.utils import logging +from areal.v2.training_service.router.config import RouterConfig +from areal.v2.training_service.router.state import ModelRegistry httpx = importlib.import_module("httpx") fastapi = importlib.import_module("fastapi") diff --git a/areal/experimental/training_service/router/config.py b/areal/v2/training_service/router/config.py similarity index 100% rename from areal/experimental/training_service/router/config.py rename to areal/v2/training_service/router/config.py diff --git a/areal/experimental/training_service/router/state.py b/areal/v2/training_service/router/state.py similarity index 100% rename from areal/experimental/training_service/router/state.py rename to areal/v2/training_service/router/state.py diff --git a/areal/v2/training_service/worker/__init__.py b/areal/v2/training_service/worker/__init__.py new file mode 100644 index 0000000000..9881313609 --- /dev/null +++ b/areal/v2/training_service/worker/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/areal/experimental/training_service/worker/__main__.py b/areal/v2/training_service/worker/__main__.py similarity index 88% rename from areal/experimental/training_service/worker/__main__.py rename to areal/v2/training_service/worker/__main__.py index 4f253ce62e..6bff8cd768 100644 --- a/areal/experimental/training_service/worker/__main__.py +++ b/areal/v2/training_service/worker/__main__.py @@ -24,9 +24,9 @@ def main(): ) args, _ = parser.parse_known_args() - from areal.experimental.training_service.worker.app import create_app - from areal.experimental.training_service.worker.config import TrainWorkerConfig from areal.infra.utils.http import validate_admin_api_key + from areal.v2.training_service.worker.app import create_app + from areal.v2.training_service.worker.config import TrainWorkerConfig validate_admin_api_key(args.host, args.admin_api_key) diff --git a/areal/experimental/training_service/worker/app.py b/areal/v2/training_service/worker/app.py similarity index 96% rename from areal/experimental/training_service/worker/app.py rename to areal/v2/training_service/worker/app.py index ba5ac5bd45..ea812ebb36 100644 --- a/areal/experimental/training_service/worker/app.py +++ b/areal/v2/training_service/worker/app.py @@ -11,13 +11,13 @@ from typing import Any from areal.api import TrainEngine -from areal.experimental.training_service.worker.awex import create_awex_blueprint -from areal.experimental.training_service.worker.config import TrainWorkerConfig -from areal.experimental.training_service.worker.engine import create_engine_module from areal.infra.platforms import current_platform from areal.infra.rpc.serialization import deserialize_value, serialize_value from areal.utils import logging from areal.utils.data import broadcast_tensor_container, tensor_container_to +from areal.v2.training_service.worker.awex import create_awex_blueprint +from areal.v2.training_service.worker.config import TrainWorkerConfig +from areal.v2.training_service.worker.engine import create_engine_module logger = logging.getLogger("TrainWorker") diff --git a/areal/experimental/training_service/worker/awex.py b/areal/v2/training_service/worker/awex.py similarity index 97% rename from areal/experimental/training_service/worker/awex.py rename to areal/v2/training_service/worker/awex.py index e206810b83..1721bdfde5 100644 --- a/areal/experimental/training_service/worker/awex.py +++ b/areal/v2/training_service/worker/awex.py @@ -195,8 +195,8 @@ def action(): def _create_training_adapter(engine): from areal.engine.fsdp_engine import FSDPEngine from areal.engine.megatron_engine import MegatronEngine - from areal.experimental.weight_update.awex.fsdp_adapter import AwexFSDPAdapter - from areal.experimental.weight_update.awex.megatron_adapter import ( + from areal.v2.weight_update.awex.fsdp_adapter import AwexFSDPAdapter + from areal.v2.weight_update.awex.megatron_adapter import ( AwexMegatronAdapter, ) diff --git a/areal/experimental/training_service/worker/config.py b/areal/v2/training_service/worker/config.py similarity index 100% rename from areal/experimental/training_service/worker/config.py rename to areal/v2/training_service/worker/config.py diff --git a/areal/experimental/training_service/worker/engine.py b/areal/v2/training_service/worker/engine.py similarity index 100% rename from areal/experimental/training_service/worker/engine.py rename to areal/v2/training_service/worker/engine.py diff --git a/areal/experimental/weight_update/__init__.py b/areal/v2/weight_update/__init__.py similarity index 81% rename from areal/experimental/weight_update/__init__.py rename to areal/v2/weight_update/__init__.py index 14fd6c00a5..b71177ffa7 100644 --- a/areal/experimental/weight_update/__init__.py +++ b/areal/v2/weight_update/__init__.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 """Weight update protocol adapters for training and inference.""" -from areal.experimental.weight_update.controller import ( +from areal.v2.weight_update.controller import ( WeightUpdateController, WeightUpdateControllerConfig, ) diff --git a/areal/experimental/weight_update/awex/__init__.py b/areal/v2/weight_update/awex/__init__.py similarity index 80% rename from areal/experimental/weight_update/awex/__init__.py rename to areal/v2/weight_update/awex/__init__.py index 4c0a78b1b4..e48c3ee343 100644 --- a/areal/experimental/weight_update/awex/__init__.py +++ b/areal/v2/weight_update/awex/__init__.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio +import os from typing import Any import aiohttp # pyright: ignore[reportMissingImports] @@ -13,6 +14,16 @@ logger = logging.getLogger("AwexHTTP") +def awex_wu_use_group() -> bool: + """Resolve whether ``batch_send_recv`` should use ``use_group=True``. + + Why: on some hardware/driver combinations, ``torch.distributed.batch_isend_irecv`` + can hang during weight update. ``AWEX_WU_USE_GROUP=0`` lets the caller fall back + to per-op send/recv to bypass the hang. Defaults to ``1`` (True). + """ + return bool(int(os.getenv("AWEX_WU_USE_GROUP", "1"))) + + async def _fetch_kv_metadata( kv_store_url: str, pair_name: str, diff --git a/areal/experimental/weight_update/awex/fsdp_adapter.py b/areal/v2/weight_update/awex/fsdp_adapter.py similarity index 97% rename from areal/experimental/weight_update/awex/fsdp_adapter.py rename to areal/v2/weight_update/awex/fsdp_adapter.py index 9c061bacb8..2f33b60e63 100644 --- a/areal/experimental/weight_update/awex/fsdp_adapter.py +++ b/areal/v2/weight_update/awex/fsdp_adapter.py @@ -19,15 +19,18 @@ from torch.distributed.tensor.placement_types import Shard from areal.engine.core.model import is_qwen_vl_model -from areal.experimental.weight_update.awex import fetch_kv_metadata -from areal.experimental.weight_update.nccl_group import ( +from areal.utils import logging +from areal.v2.weight_update.awex import ( + awex_wu_use_group, + fetch_kv_metadata, +) +from areal.v2.weight_update.nccl_group import ( init_weights_update_group, setup_batch_isend_irecv, ) -from areal.experimental.weight_update.training_adapter import ( +from areal.v2.weight_update.training_adapter import ( AwexTrainingAdapter, ) -from areal.utils import logging if TYPE_CHECKING: from areal.engine.fsdp_engine import FSDPEngine @@ -174,7 +177,12 @@ def execute_weight_update(self, version: int) -> None: self._weights_update_group, copy_rank=self._transfer_rank, ) - batch_send_recv(send_ops=send_ops, recv_ops=[], blocking=True) + batch_send_recv( + send_ops=send_ops, + recv_ops=[], + blocking=True, + use_group=awex_wu_use_group(), + ) torch.distributed.barrier(group=self._weights_update_group) def batch_isend_irecv(self, **kwargs) -> None: diff --git a/areal/experimental/weight_update/awex/megatron_adapter.py b/areal/v2/weight_update/awex/megatron_adapter.py similarity index 98% rename from areal/experimental/weight_update/awex/megatron_adapter.py rename to areal/v2/weight_update/awex/megatron_adapter.py index 6002a7fa8b..85306ff9c3 100644 --- a/areal/experimental/weight_update/awex/megatron_adapter.py +++ b/areal/v2/weight_update/awex/megatron_adapter.py @@ -24,15 +24,18 @@ group_tensors_by_shape_and_dtype, ) -from areal.experimental.weight_update.awex import fetch_kv_metadata -from areal.experimental.weight_update.nccl_group import ( +from areal.utils import logging +from areal.v2.weight_update.awex import ( + awex_wu_use_group, + fetch_kv_metadata, +) +from areal.v2.weight_update.nccl_group import ( init_weights_update_group, setup_batch_isend_irecv, ) -from areal.experimental.weight_update.training_adapter import ( +from areal.v2.weight_update.training_adapter import ( AwexTrainingAdapter, ) -from areal.utils import logging if TYPE_CHECKING: from areal.engine.megatron_engine import MegatronEngine @@ -198,7 +201,12 @@ def execute_weight_update(self, version: int) -> None: self._weights_update_group, copy_rank=self._transfer_rank, ) - batch_send_recv(send_ops=send_ops, recv_ops=[], blocking=True) + batch_send_recv( + send_ops=send_ops, + recv_ops=[], + blocking=True, + use_group=awex_wu_use_group(), + ) dist.barrier(group=self._weights_update_group) def batch_isend_irecv(self, **kwargs) -> None: diff --git a/areal/experimental/weight_update/awex/sglang_adapter.py b/areal/v2/weight_update/awex/sglang_adapter.py similarity index 98% rename from areal/experimental/weight_update/awex/sglang_adapter.py rename to areal/v2/weight_update/awex/sglang_adapter.py index 6194031806..d70b5d2cd0 100644 --- a/areal/experimental/weight_update/awex/sglang_adapter.py +++ b/areal/v2/weight_update/awex/sglang_adapter.py @@ -30,15 +30,18 @@ reconstruct_tensors_from_groups, ) -from areal.experimental.weight_update.awex import fetch_kv_metadata -from areal.experimental.weight_update.inference_adapter import ( +from areal.utils import logging +from areal.v2.weight_update.awex import ( + awex_wu_use_group, + fetch_kv_metadata, +) +from areal.v2.weight_update.inference_adapter import ( AwexInferenceAdapter, ) -from areal.experimental.weight_update.nccl_group import ( +from areal.v2.weight_update.nccl_group import ( init_weights_update_group, setup_batch_isend_irecv, ) -from areal.utils import logging logger = logging.getLogger("AwexSGLangAdapter") @@ -409,7 +412,12 @@ def execute_weight_update(self, version: int) -> None: self._transfer_plan, self._weights_update_group, ) - batch_send_recv(send_ops=[], recv_ops=recv_ops, blocking=True) + batch_send_recv( + send_ops=[], + recv_ops=recv_ops, + blocking=True, + use_group=awex_wu_use_group(), + ) for original, contiguous in non_contiguous_pairs: original.copy_(contiguous) diff --git a/areal/experimental/weight_update/controller/__init__.py b/areal/v2/weight_update/controller/__init__.py similarity index 58% rename from areal/experimental/weight_update/controller/__init__.py rename to areal/v2/weight_update/controller/__init__.py index 09f0587d84..fd7bd8c7a0 100644 --- a/areal/experimental/weight_update/controller/__init__.py +++ b/areal/v2/weight_update/controller/__init__.py @@ -1,9 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -from areal.experimental.weight_update.controller.config import ( +from areal.v2.weight_update.controller.config import ( WeightUpdateControllerConfig, ) -from areal.experimental.weight_update.controller.controller import ( +from areal.v2.weight_update.controller.controller import ( WeightUpdateController, ) diff --git a/areal/experimental/weight_update/controller/config.py b/areal/v2/weight_update/controller/config.py similarity index 100% rename from areal/experimental/weight_update/controller/config.py rename to areal/v2/weight_update/controller/config.py diff --git a/areal/experimental/weight_update/controller/controller.py b/areal/v2/weight_update/controller/controller.py similarity index 94% rename from areal/experimental/weight_update/controller/controller.py rename to areal/v2/weight_update/controller/controller.py index 558ea7fb9e..457cc5d273 100644 --- a/areal/experimental/weight_update/controller/controller.py +++ b/areal/v2/weight_update/controller/controller.py @@ -9,13 +9,13 @@ import httpx -from areal.experimental.weight_update.controller.config import ( - WeightUpdateControllerConfig, -) -from areal.experimental.weight_update.gateway.config import WeightUpdateResult from areal.infra.utils.proc import kill_process_tree from areal.utils import logging from areal.utils.network import find_free_ports +from areal.v2.weight_update.controller.config import ( + WeightUpdateControllerConfig, +) +from areal.v2.weight_update.gateway.config import WeightUpdateResult logger = logging.getLogger("WeightUpdateController") @@ -47,7 +47,7 @@ def initialize(self) -> None: cmd = [ sys.executable, "-m", - "areal.experimental.weight_update.gateway", + "areal.v2.weight_update.gateway", "--host", cfg.host, "--port", @@ -112,6 +112,7 @@ def connect( save_path: str = "", use_lora: bool = False, lora_name: str = "", + lora_keep_versions: int = 0, colocate: bool = False, nccl_master_addr: str = "", nccl_master_port: int = 0, @@ -125,6 +126,7 @@ def connect( "save_path": save_path, "use_lora": use_lora, "lora_name": lora_name, + "lora_keep_versions": lora_keep_versions, "colocate": colocate, "nccl_master_addr": nccl_master_addr, "nccl_master_port": nccl_master_port, @@ -136,7 +138,11 @@ def connect( ) resp.raise_for_status() logger.info( - "Connected pair '%s' (mode=%s, colocate=%s)", pair_name, mode, colocate + "Connected pair '%s' (mode=%s, colocate=%s, use_lora=%s)", + pair_name, + mode, + colocate, + use_lora, ) def update_weights(self, version: int) -> WeightUpdateResult: diff --git a/areal/v2/weight_update/gateway/__init__.py b/areal/v2/weight_update/gateway/__init__.py new file mode 100644 index 0000000000..9881313609 --- /dev/null +++ b/areal/v2/weight_update/gateway/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/areal/experimental/weight_update/gateway/__main__.py b/areal/v2/weight_update/gateway/__main__.py similarity index 88% rename from areal/experimental/weight_update/gateway/__main__.py rename to areal/v2/weight_update/gateway/__main__.py index 513102e8b6..585858537b 100644 --- a/areal/experimental/weight_update/gateway/__main__.py +++ b/areal/v2/weight_update/gateway/__main__.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 -"""CLI entrypoint: ``python -m areal.experimental.weight_update.gateway``.""" +"""CLI entrypoint: ``python -m areal.v2.weight_update.gateway``.""" from __future__ import annotations @@ -36,13 +36,13 @@ def main(): ) args, _ = parser.parse_known_args() - from areal.experimental.weight_update.gateway.app import create_app - from areal.experimental.weight_update.gateway.config import WeightUpdateConfig from areal.infra.utils.http import ( get_default_uvicorn_kwargs, validate_admin_api_key, ) from areal.utils.logging import suppress_http_loggers + from areal.v2.weight_update.gateway.app import create_app + from areal.v2.weight_update.gateway.config import WeightUpdateConfig validate_admin_api_key(args.host, args.admin_api_key) diff --git a/areal/experimental/weight_update/gateway/app.py b/areal/v2/weight_update/gateway/app.py similarity index 90% rename from areal/experimental/weight_update/gateway/app.py rename to areal/v2/weight_update/gateway/app.py index 0a93b7fd19..a381b8211c 100644 --- a/areal/experimental/weight_update/gateway/app.py +++ b/areal/v2/weight_update/gateway/app.py @@ -13,17 +13,17 @@ from fastapi.responses import JSONResponse # pyright: ignore[reportMissingImports] from pydantic import BaseModel # pyright: ignore[reportMissingImports] -from areal.experimental.weight_update.gateway.auth import require_admin_key -from areal.experimental.weight_update.gateway.config import ( +from areal.infra.utils.http import async_http_retry +from areal.utils import logging +from areal.utils.network import find_free_ports +from areal.v2.weight_update.gateway.auth import require_admin_key +from areal.v2.weight_update.gateway.config import ( PairInfo, WeightUpdateConfig, WeightUpdateResult, ) -from areal.experimental.weight_update.gateway.kv_store import WeightMetaStore -from areal.experimental.weight_update.gateway.pair_registry import PairRegistry -from areal.infra.utils.http import async_http_retry -from areal.utils import logging -from areal.utils.network import find_free_ports +from areal.v2.weight_update.gateway.kv_store import WeightMetaStore +from areal.v2.weight_update.gateway.pair_registry import PairRegistry logger = logging.getLogger("WeightUpdateGateway") @@ -38,6 +38,7 @@ class ConnectRequest(BaseModel): save_path: str = "" use_lora: bool = False lora_name: str = "" + lora_keep_versions: int = 0 colocate: bool = False @@ -229,13 +230,35 @@ async def connect(request: Request, body: ConnectRequest) -> ConnectResponse: save_path=body.save_path, use_lora=body.use_lora, lora_name=body.lora_name, + lora_keep_versions=body.lora_keep_versions, ) registry.register(pair_info) logger.info( - "Connected disk pair '%s' (save_path=%s)", pair_name, body.save_path + "Connected disk pair '%s' (save_path=%s, use_lora=%s, " + "lora_keep_versions=%d)", + pair_name, + body.save_path, + body.use_lora, + body.lora_keep_versions, ) return ConnectResponse(pair_name=pair_name) + # awex mode -- LoRA is unsupported because the NCCL P2P transfer plan + # assumes train/infer parameter names match the HF layout, but PEFT + # exposes ``base_model.model.*.{base_layer,lora_A,lora_B}.weight`` on + # the train side. Fail fast with an actionable error rather than + # bubbling up a cryptic TransferPlanBuilder key-mismatch at init. + if body.use_lora: + return JSONResponse( + status_code=400, + content={ + "error": ( + "awex weight update does not support LoRA; set " + "actor.weight_update_mode=disk in your config." + ) + }, + ) + session = request.app.state.http_session init_timeout_s = config.init_timeout_s @@ -633,6 +656,35 @@ async def _disk_transfer_weights( for url in pair_info.inference_worker_urls ] ) + # Unload the version that fell outside the retention window so + # sglang does not accumulate one adapter per train step (leaks + # VRAM and eventually hangs). Best-effort: the stale adapter may + # already have been evicted or never loaded on this worker. + keep = pair_info.lora_keep_versions + if keep > 0 and version - keep >= 0: + stale_name = get_versioned_lora_name( + pair_info.lora_name, version - keep + ) + + async def _unload(url: str) -> None: + try: + await _post_json( + session, + f"{url}/unload_lora_adapter", + timeout_s, + json_data={"lora_name": stale_name}, + ) + except Exception as e: + logger.warning( + "unload_lora_adapter(%s) on %s failed (best-effort): %s", + stale_name, + url, + e, + ) + + await asyncio.gather( + *[_unload(url) for url in pair_info.inference_worker_urls] + ) else: await asyncio.gather( *[ diff --git a/areal/experimental/weight_update/gateway/auth.py b/areal/v2/weight_update/gateway/auth.py similarity index 100% rename from areal/experimental/weight_update/gateway/auth.py rename to areal/v2/weight_update/gateway/auth.py diff --git a/areal/experimental/weight_update/gateway/config.py b/areal/v2/weight_update/gateway/config.py similarity index 86% rename from areal/experimental/weight_update/gateway/config.py rename to areal/v2/weight_update/gateway/config.py index db564d8c77..34c40b4f59 100644 --- a/areal/experimental/weight_update/gateway/config.py +++ b/areal/v2/weight_update/gateway/config.py @@ -59,6 +59,10 @@ class PairInfo: save_path: str = "" use_lora: bool = False lora_name: str = "" + # Number of recent LoRA adapter versions to keep loaded on each inference + # worker. Older versions are unloaded after a successful load_lora_adapter + # to bound VRAM and avoid the adapter-accumulation hang. 0 disables cleanup. + lora_keep_versions: int = 0 # Colocated mode (training and inference share GPUs) colocate: bool = False diff --git a/areal/experimental/weight_update/gateway/kv_store.py b/areal/v2/weight_update/gateway/kv_store.py similarity index 100% rename from areal/experimental/weight_update/gateway/kv_store.py rename to areal/v2/weight_update/gateway/kv_store.py diff --git a/areal/experimental/weight_update/gateway/pair_registry.py b/areal/v2/weight_update/gateway/pair_registry.py similarity index 94% rename from areal/experimental/weight_update/gateway/pair_registry.py rename to areal/v2/weight_update/gateway/pair_registry.py index 2b23b2497c..bbec35a565 100644 --- a/areal/experimental/weight_update/gateway/pair_registry.py +++ b/areal/v2/weight_update/gateway/pair_registry.py @@ -3,8 +3,8 @@ import threading -from areal.experimental.weight_update.gateway.config import PairInfo from areal.utils import logging +from areal.v2.weight_update.gateway.config import PairInfo logger = logging.getLogger("PairRegistry") diff --git a/areal/experimental/weight_update/inference_adapter.py b/areal/v2/weight_update/inference_adapter.py similarity index 100% rename from areal/experimental/weight_update/inference_adapter.py rename to areal/v2/weight_update/inference_adapter.py diff --git a/areal/experimental/weight_update/nccl_group.py b/areal/v2/weight_update/nccl_group.py similarity index 100% rename from areal/experimental/weight_update/nccl_group.py rename to areal/v2/weight_update/nccl_group.py diff --git a/areal/experimental/weight_update/training_adapter.py b/areal/v2/weight_update/training_adapter.py similarity index 100% rename from areal/experimental/weight_update/training_adapter.py rename to areal/v2/weight_update/training_adapter.py diff --git a/areal/workflow/rlvr.py b/areal/workflow/rlvr.py index aa8e28a9a6..365d0ca479 100644 --- a/areal/workflow/rlvr.py +++ b/areal/workflow/rlvr.py @@ -1,22 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + import uuid from collections.abc import Callable -from typing import Any +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any import torch from transformers import PreTrainedTokenizerFast -from areal import workflow_context from areal.api import ( AsyncRewardWrapper, - InferenceEngine, - ModelRequest, - ModelResponse, + RewardResult, RolloutWorkflow, + normalize_reward_result, ) -from areal.api.cli_args import GenerationHyperparameters -from areal.utils import logging, stats_tracker +from areal.utils import logging from areal.utils.dynamic_import import import_from_string from areal.utils.hf_utils import apply_chat_template from areal.utils.perf_tracer import ( @@ -25,9 +25,22 @@ trace_session, ) +if TYPE_CHECKING: + from areal.api.engine_api import InferenceEngine + from areal.api.io_struct import ModelResponse + logger = logging.getLogger("RLVRWorkflow") +@dataclass +class _WorkflowModelRequest: + rid: str + input_ids: list[int] + gconfig: Any + tokenizer: PreTrainedTokenizerFast | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + def default_get_input_ids_fn( data: Any, tokenizer: PreTrainedTokenizerFast, @@ -52,7 +65,7 @@ class RLVRWorkflow(RolloutWorkflow): def __init__( self, reward_fn: Callable[..., Any] | str, - gconfig: GenerationHyperparameters, + gconfig: Any, tokenizer: PreTrainedTokenizerFast | str, enable_thinking: bool = False, get_input_ids_fn: Callable[[Any, PreTrainedTokenizerFast, bool], list[int]] @@ -86,7 +99,7 @@ async def _compute_rewards( resp: ModelResponse, prompt_str: str, task_data: dict[str, Any], - ) -> float: + ) -> RewardResult: """Decode completion and compute reward. Traces reward phase execution for SessionTracer. Decodes output tokens @@ -94,8 +107,8 @@ async def _compute_rewards( Returns ------- - float - Reward value. + RewardResult + Reward payload. """ completions_str = self.tokenizer.decode(resp.output_tokens) reward = await self.async_reward_fn( @@ -106,16 +119,16 @@ async def _compute_rewards( **task_data, ) - return reward + return normalize_reward_result(reward) @session_context() async def _collect_samples( self, - engine: InferenceEngine, - req: ModelRequest, + engine: "InferenceEngine", + req: _WorkflowModelRequest, prompt_str: str, task_data: dict[str, Any], - ) -> tuple[ModelResponse, float]: + ) -> tuple["ModelResponse", RewardResult]: """Generate one sample and compute its reward. Registers a new session for this sample, calls engine.agenerate, @@ -124,20 +137,69 @@ async def _collect_samples( Returns ------- - tuple[ModelResponse, float] - Model response and reward value. + tuple[ModelResponse, RewardResult] + Model response and reward payload. """ async with atrace_session_phase("generate"): resp = await engine.agenerate(req) reward = await self._compute_rewards(resp, prompt_str, task_data) - stats_tracker.get(workflow_context.stat_scope()).scalar(reward=reward) + from areal.infra import workflow_context + from areal.utils import stats_tracker + + tracker = stats_tracker.get(workflow_context.stat_scope()) + tracker.scalar(reward=reward.final_reward) + if reward.step_rewards is not None: + tracker.scalar( + step_reward_count=len(reward.step_rewards), + step_reward_sum=float(sum(reward.step_rewards)), + ) return resp, reward + def _build_stepwise_reward_tensors( + self, + reward: RewardResult, + resp: "ModelResponse", + ) -> tuple[torch.Tensor, torch.Tensor]: + """Convert stepwise rewards into completion-aligned tensors. + + Returns tensors of shape ``[seqlen]`` aligned to the full prompt + + completion sequence. Reward values are injected at the completion token + positions indicated by ``reward.step_ends``. + """ + + seq_len = resp.input_len + resp.output_len + step_rewards = torch.zeros(seq_len, dtype=torch.float32) + step_reward_mask = torch.zeros(seq_len, dtype=torch.bool) + + if reward.step_rewards is None and reward.step_ends is None: + return step_rewards, step_reward_mask + if reward.step_rewards is None or reward.step_ends is None: + raise ValueError( + "step_rewards and step_ends must either both be provided or both be None" + ) + if len(reward.step_rewards) != len(reward.step_ends): + raise ValueError( + "step_rewards and step_ends must have the same length, got " + f"{len(reward.step_rewards)} and {len(reward.step_ends)}" + ) + + completion_len = resp.output_len + for step_reward, step_end in zip(reward.step_rewards, reward.step_ends): + if step_end <= 0 or step_end > completion_len: + raise ValueError( + f"Invalid step_end={step_end}; expected 1 <= step_end <= {completion_len}" + ) + seq_index = resp.input_len + step_end - 1 + step_rewards[seq_index] += float(step_reward) + step_reward_mask[seq_index] = True + + return step_rewards, step_reward_mask + async def arun_episode( - self, engine: InferenceEngine, data: dict[str, Any] + self, engine: "InferenceEngine", data: dict[str, Any] ) -> dict[str, torch.Tensor]: # NOTE: load reward function dynamically if given as string if isinstance(self.reward_fn, str): @@ -149,7 +211,7 @@ async def arun_episode( self.tokenizer, self.enable_thinking, ) - req = ModelRequest( + req = _WorkflowModelRequest( rid=uuid.uuid4().hex, input_ids=input_ids, gconfig=self.gconfig.new(n_samples=1), @@ -166,6 +228,9 @@ async def arun_episode( logprobs = [0.0] * resp.input_len + resp.output_logprobs loss_mask = [0] * resp.input_len + [1] * resp.output_len versions = [-1] * resp.input_len + resp.output_versions + step_rewards, step_reward_mask = self._build_stepwise_reward_tensors( + reward, resp + ) res = { "input_ids": torch.tensor(seq, dtype=torch.int32), @@ -173,6 +238,8 @@ async def arun_episode( "logprobs": torch.tensor(logprobs, dtype=torch.float32), "versions": torch.tensor(versions, dtype=torch.int32), "attention_mask": torch.ones(len(seq), dtype=torch.bool), - "rewards": torch.tensor(reward, dtype=torch.float32), + "rewards": torch.tensor(reward.final_reward, dtype=torch.float32), + "step_rewards": step_rewards, + "step_reward_mask": step_reward_mask, } return {k: v.unsqueeze(0) for k, v in res.items()} diff --git a/assets/figures/wechat_qrcode.png b/assets/figures/wechat_qrcode.png index 885cc31dd9..01246f26dc 100644 Binary files a/assets/figures/wechat_qrcode.png and b/assets/figures/wechat_qrcode.png differ diff --git a/docs/en/algorithms/grpo_series.md b/docs/en/algorithms/grpo_series.md index 0eb3dc1f8d..4943b530c3 100644 --- a/docs/en/algorithms/grpo_series.md +++ b/docs/en/algorithms/grpo_series.md @@ -18,6 +18,11 @@ training, including: - **SAPO**: [Paper](https://arxiv.org/abs/2511.20347) - **GSPO** (Qwen3): [Paper](https://arxiv.org/abs/2507.18071), [Blog](https://qwenlm.github.io/blog/gspo/) +- **IcePop**: [Blog](https://ringtech.notion.site/icepop) — Importance-ratio-based token masking (composable with other RL algorithms) +- **KPop**: [Blog](https://ringtech.notion.site/kpop) — Bidirectional binary KL divergence token masking (composable with other RL algorithms) + +IcePop and KPop are token masking strategies that can be composed with any RL algorithm +listed above. These algorithms share the same base objective but differ in their normalization strategies, clipping mechanisms, importance sampling levels, etc. By adjusting a few @@ -35,7 +40,7 @@ configuration YAML file. | **slurm** | `python3 examples/math/gsm8k_rl.py --config examples/math/gsm8k_.yaml scheduler.type=slurm` | Replace `` with: `ppo`, `grpo`, `drgrpo`, `liteppo`, `rloo`, `gspo`, -`dapo_dynamic_bs`, or `sapo`. +`dapo_dynamic_bs`, `sapo`, `icepop`, or `kpop`. ### Switching Algorithms via CLI Overrides @@ -148,8 +153,10 @@ parameters: | **LitePPO** | `group` | `batch` | `false` | `token` | - | | **RLOO** | `group` | `null` | `true` | `token` | - | | **GSPO** | `batch` | `batch` | `false` | `sequence` | - | -| **DAPO** | `batch` | `batch` | `false` | `token` | asymmetric clip, dynamic sampling | -| **SAPO** | `batch` | `batch` | `false` | `token` | `use_sapo_loss=true` | +| **DAPO** | `batch` | `batch` | `false` | `token` | asymmetric clip, dynamic sampling | +| **SAPO** | `batch` | `batch` | `false` | `token` | `use_sapo_loss=true` | +| **IcePop** | `batch` | `batch` | `false` | `token` | `rejection_sampling.metric=ratio` | +| **KPop** | `batch` | `batch` | `false` | `token` | `rejection_sampling.metric=binary_kl` | **Note**: The "GRPO" row reflects the original DeepSeekMath formulation. AReaL's default GRPO config uses these settings but with length normalization already removed (see AReaL @@ -298,6 +305,62 @@ and then filter them. The following option controls batch sizing behavior: | ------------ | ---- | ------- | --------------------------- | | `dynamic_bs` | bool | `false` | Enable dynamic batch sizing | +### IcePop + +IcePop masks tokens whose importance ratio $r_{i,t} = \frac{\pi_\theta(o_{i,t} \mid q, o_{i, \phi$ are masked (here $\phi$ corresponds to `actor.rejection_sampling.upper`). + +It is implemented via the `rejection_sampling` config with `metric=binary_kl`: + +```yaml +actor: + use_decoupled_loss: true + rejection_sampling: + level: token + action: mask + metric: binary_kl + upper: 2.0 +``` + +| Parameter | Type | Default | Description | +| ------------------------------------ | ----------- | ------- | ------------------------------------ | +| `actor.rejection_sampling.metric` | str | - | Set to `binary_kl` for KPop | +| `actor.rejection_sampling.upper` | float | `2.0` | KL divergence threshold ($\phi$) | + +**Note:** KPop only supports `action=mask` (not `clamp`), and `lower` is not used with `binary_kl`. KPop requires `actor.use_decoupled_loss=true`, otherwise `rejection_sampling` has no effect. + +See `examples/math/gsm8k_kpop.yaml` for a complete configuration example. + ## Core Concepts **Rewards**: AReaL assumes outcome-based rewards. Each trajectory, which may consist of diff --git a/docs/en/cli_reference.md b/docs/en/cli_reference.md index ec1272c6f2..5e4b9ca5d6 100644 --- a/docs/en/cli_reference.md +++ b/docs/en/cli_reference.md @@ -119,6 +119,7 @@ Base configuration class for all experiment types with common settings. | `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | | `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | | `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | +| `post_exit_hook` | string | `""` | Shell command run after launcher shutdown. LOG_DIR is injected; failures are logged and ignored. | (section-grpo)= @@ -149,6 +150,7 @@ A dummy place holder of GRPO config for backward compatibility. | `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | | `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | | `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | +| `post_exit_hook` | string | `""` | Shell command run after launcher shutdown. LOG_DIR is injected; failures are logged and ignored. | | `gconfig` | [`GenerationHyperparameters`](section-generation-hyperparameters) | **Required** | - | | `eval_gconfig` | [`GenerationHyperparameters`](section-generation-hyperparameters) \| None | `None` | Generation hyperparameters for evaluation. If None, use gconfig. | | `rollout` | [`InferenceEngineConfig`](section-inference-engine) | **Required** | - | @@ -187,6 +189,7 @@ Configuration for Proximal Policy Optimization (PPO) reinforcement learning expe | `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | | `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | | `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | +| `post_exit_hook` | string | `""` | Shell command run after launcher shutdown. LOG_DIR is injected; failures are logged and ignored. | | `gconfig` | [`GenerationHyperparameters`](section-generation-hyperparameters) | **Required** | - | | `eval_gconfig` | [`GenerationHyperparameters`](section-generation-hyperparameters) \| None | `None` | Generation hyperparameters for evaluation. If None, use gconfig. | | `rollout` | [`InferenceEngineConfig`](section-inference-engine) | **Required** | - | @@ -225,6 +228,7 @@ Configuration for Reward Model (RW) training experiments. | `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | | `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | | `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | +| `post_exit_hook` | string | `""` | Shell command run after launcher shutdown. LOG_DIR is injected; failures are logged and ignored. | | `actor` | [`TrainEngineConfig`](section-train-engine) | **Required** | - | (section-sft)= @@ -256,6 +260,7 @@ Configuration for Supervised Fine-Tuning (SFT) experiments. | `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | | `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | | `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | +| `post_exit_hook` | string | `""` | Shell command run after launcher shutdown. LOG_DIR is injected; failures are logged and ignored. | | `actor` | [`TrainEngineConfig`](section-train-engine) | **Required** | - | (section-fsdp-engine)= @@ -401,11 +406,12 @@ Configuration for PPO actor model, a subclass of a TrainEngine. | `use_sapo_loss` | boolean | `False` | Use SAPO loss (mutually exclusive with PPO clipping) | | `sapo_tau_pos` | float | `1.0` | SAPO temperature for positive advantages | | `sapo_tau_neg` | float | `1.05` | SAPO temperature for negative advantages | +| `use_cispo_loss` | boolean | `False` | Use CISPO loss: clip the importance-sampling weight under stop-gradient and keep gradient on every token's log pi (MiniMax-M1 Eq. 4-5). Mutually exclusive with SAPO. Token-level only. Requires eps_clip_higher > 0; recommended eps_clip=1.0 (single-sided, lower bound 0) with eps_clip_higher=4.0. | | `recompute_logprob` | boolean | `False` | Recompute log probability and replace the log probability returned by inference. | | `use_decoupled_loss` | boolean | `False` | Use the decoupled loss. Implicitly enables recompute_logprob. | | `rejection_sampling` | [`RejectionSamplingConfig`](section-rejection-sampling) \| None | `None` | Rejection sampling configuration for filtering stale samples. None disables filtering (equivalent to old behave_imp_weight_mode='disabled'). Only effective when use_decoupled_loss=True. | | `importance_sampling_level` | string | `"token"` | Level at which to compute importance sampling ratios. 'token': per-token ratios (standard PPO). 'sequence': sequence-level geometric mean of per-token ratios (GSPO). **Choices:** `token`, `sequence` | -| `prox_logp_method` | string | `"recompute"` | Method for computing proximal policy log-probabilities in decoupled PPO. Only effective when use_decoupled_loss=True. Options: 'recompute' (default): Standard decoupled PPO, recompute proximal policy via forward pass. 'loglinear': Use log-linear interpolation to approximate proximal policy (skip forward pass). 'metrics': Like 'recompute', but also compute approximation metrics for evaluation. **Choices:** `recompute`, `loglinear`, `metrics` | +| `prox_logp_method` | string | `"recompute"` | Method for computing proximal policy log-probabilities in decoupled PPO. Only effective when use_decoupled_loss=True. Options: 'recompute' (default): Standard decoupled PPO, recompute proximal policy via forward pass. 'loglinear': Use log-linear interpolation to approximate proximal policy (skip forward pass). 'metrics': Like 'recompute', but also compute approximation metrics for evaluation. 'reuse_train_logp': Reuse training forward-pass logprobs as the proximal logp (skip the extra forward; requires ppo_n_minibatches=1). **Choices:** `recompute`, `loglinear`, `metrics`, `reuse_train_logp` | | `log_agent_stats` | boolean | `False` | Log statistics for agent trajectories | | `log_agent_stats_keys` | list of string | **Required** | Keys for logging agent trajectory statistics | | `max_new_tokens` | integer | `1024` | Maximum number of new tokens to generate | @@ -893,6 +899,8 @@ Consolidates proxy settings (mode, parsers, export) with agent-service orchestra | `engine_max_tokens` | integer \| None | `None` | Maximum total tokens for the engine (prompt + completion). | | `turn_discount` | float | `1.0` | Discount factor for multi-turn reward propagation. | | `export_style` | string | `"individual"` | Export style: 'individual' (all interactions) or 'concat' (leaf nodes only). The 'individual' style exports each interaction (input-output-reward) step separately, and treats them as independent samples to train the model. The 'concat' style exports only the final concatenated trajectory from the root. It is only suitable for linear conversation histories without token mismatching (whether valid depends on the tokenizer). **Choices:** `individual`, `concat` | +| `message_preprocessors` | list of string | **Required** | List of message preprocessor class paths applied, in order, to Anthropic-compatible `/v1/messages` requests after translating them to OpenAI-compatible requests. Native OpenAI `/chat/completions` and `/responses` requests are not preprocessed. Each entry is a dotted import path to a callable class. | +| `prefix_matcher` | string \| None | `None` | Dotted import path to a custom prefix matcher function for InteractionCache parent-child matching. The function must accept two list\[dict\] arguments (candidate prefix, full messages) and return bool. When None, exact element-wise equality is used. | | `subproc_max_workers` | integer | `4` | Maximum number of worker processes for subprocess mode execution pool. | | `session_timeout_seconds` | integer | `3600` | Session timeout in seconds. Sessions inactive longer than this will be garbage collected. | | `set_reward_finish_timeout` | float | `0.0` | Timeout in seconds to wait for additional reward updates before finalizing a session. | @@ -964,6 +972,7 @@ Configuration for Direct Preference Optimization (DPO) experiments. | `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | | `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | | `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | +| `post_exit_hook` | string | `""` | Shell command run after launcher shutdown. LOG_DIR is injected; failures are logged and ignored. | | `actor` | [`DPOEngineConfig`](section-dpo-engine) | **Required** | - | | `ref` | [`DPOEngineConfig`](section-dpo-engine) | **Required** | - | @@ -1096,10 +1105,17 @@ Refer to Megatron-LM documentation for implementation details. | `moe_enable_deepep` | boolean | `False` | - | | `moe_token_dispatcher_type` | string | `"alltoall"` | Type of token dispatcher. Options: 'allgather','alltoall' and 'flex'. | | `moe_permute_fusion` | boolean | `False` | Fuse token rearrangement ops during token dispatching. | +| `moe_router_fusion` | boolean | `False` | Enable fusion for MoE TopK routing and aux-loss computation. Requires TransformerEngine >= 2.7.0. | +| `moe_router_bias_update_rate` | float | `0.0` | Update rate for auxiliary-loss-free MoE load balancing (DeepSeek V3 style). Controls how fast expert_bias adjusts. Default 0.0 disables bias updates; set a positive value such as 1e-3 to enable. | +| `moe_z_loss_coeff` | float \| None | `None` | Scaling coefficient for router z-loss. Complements auxiliary-loss-free load balancing for router stability. A starting value of 1e-3 is recommended. None disables z-loss. | +| `enable_fp32_lm_head` | boolean | `False` | Cast lm_head output to FP32 before loss computation for numerical stability. | +| `cross_entropy_loss_fusion` | boolean | `False` | Enable fused cross-entropy loss kernel for better performance. | | `fp8_config` | [`FP8EngineConfig`](section-fp8-engine) \| None | `None` | - | | `bridge_type` | string | `"mbridge"` | Bridge backend for MegatronEngine. Choices: 'mbridge' or 'megatron-bridge'. **Choices:** `mbridge`, `megatron-bridge` | | `use_mbridge_save` | boolean | `False` | Use mbridge's save method to save gpu memory when saving weights. | | `use_bridge_for_update_weights` | boolean | `False` | When True and bridge_type='megatron-bridge', delegate live weight sync to bridge.export_hf_weights instead of the hand-rolled convert_to_hf registry. Required for models without a registry entry (e.g. Qwen3.5). FP8 paths fall back to the registry automatically. | +| `disable_grad_buffers_cpu_backup` | boolean | `False` | When offloading with torch_memory_saver, skip CPU backup for Megatron gradient buffers (they are recomputed each step). | +| `enable_mtp` | boolean | `False` | Keep the model's Multi-Token-Prediction (MTP) head (bridge_type=megatron-bridge only). Default False drops it. | (section-memory-profiler)= @@ -1163,14 +1179,14 @@ Attributes: lower: Lower bound for filtering (optional). ``` -| Parameter | Type | Default | Description | -| --------- | ------------- | --------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `level` | string | `"token"` | Filtering granularity. 'token': per-token filtering (each token judged independently). 'sequence': per-sequence filtering (all tokens in a sequence share the same fate). When metric='ratio', both the filtering decision and the correction weight (behave_imp_weight) operate at sequence level using the geometric mean. **Choices:** `token`, `sequence` | -| `action` | string | `"mask"` | Action to take when metric exceeds threshold. 'mask': zero out loss_mask for filtered tokens/sequences (rejection, completely excludes from gradient computation). 'clamp': clamp importance weight to \[lower, upper\] bounds (truncation, tokens still participate in gradient but with bounded weight). **Choices:** `mask`, `clamp` | -| `metric` | string | `"ratio"` | Divergence metric for filtering. 'ratio': direct importance ratio π_proximal/π_behave. 'kl_k1': KL estimator k1 = log(r), forward KL unbiased estimator (can be negative). 'kl_k2': KL estimator k2 = 0.5 * (log r)^2, non-negative quadratic approximation. 'kl_k3': KL estimator k3 = r - log(r) - 1, non-negative exact forward KL estimator. **Choices:** `ratio`, `kl_k1`, `kl_k2`, `kl_k3` | -| `agg` | string | `"mean"` | Aggregation method for sequence-level filtering. Only used when level='sequence'. For 'ratio' metric, aggregation is in log space: 'sum' = exp(sum(log(r_i))), 'mean' = exp(mean(log(r_i))) = geometric mean (length-invariant, consistent with GSPO). For KL metrics, aggregation is arithmetic: 'sum' = sum(kl_i), 'mean' = mean(kl_i). 'max': max of per-token metric values (most conservative). **Choices:** `sum`, `mean`, `max` | -| `upper` | float | `5.0` | Upper bound for filtering. Tokens/sequences with metric > upper are filtered out (loss_mask zeroed). For 'ratio' metric: must be > 1.0, typical values are 2.0 or 5.0. For 'kl_k2'/'kl_k3' metrics: typical values are 0.5-2.0. | -| `lower` | float \| None | `None` | Lower bound for filtering (optional). None means no lower bound. For 'ratio' metric: typical value is 0.5 (filter out tokens where policy probability dropped significantly). Must be > 0. For 'kl_k1' metric: can be used to filter negative KL estimates. | +| Parameter | Type | Default | Description | +| --------- | ------------- | --------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `level` | string | `"token"` | Filtering granularity. 'token': per-token filtering (each token judged independently). 'sequence': per-sequence filtering (all tokens in a sequence share the same fate). When metric='ratio', both the filtering decision and the correction weight (behave_imp_weight) operate at sequence level using the geometric mean. **Choices:** `token`, `sequence` | +| `action` | string | `"mask"` | Action to take when metric exceeds threshold. 'mask': zero out loss_mask for filtered tokens/sequences (rejection, completely excludes from gradient computation). 'clamp': clamp importance weight to \[lower, upper\] bounds (truncation, tokens still participate in gradient but with bounded weight). **Choices:** `mask`, `clamp` | +| `metric` | string | `"ratio"` | Divergence metric for filtering. 'ratio': direct importance ratio π_proximal/π_behave. 'kl_k1': KL estimator k1 = log(r), forward KL unbiased estimator (can be negative). 'kl_k2': KL estimator k2 = 0.5 * (log r)^2, non-negative quadratic approximation. 'kl_k3': KL estimator k3 = r - log(r) - 1, non-negative exact forward KL estimator. 'binary_kl': KPop (symmetric binary KL divergence) — masks tokens where either KL(proximal | +| `agg` | string | `"mean"` | Aggregation method for sequence-level filtering. Only used when level='sequence'. For 'ratio' metric, aggregation is in log space: 'sum' = exp(sum(log(r_i))), 'mean' = exp(mean(log(r_i))) = geometric mean (length-invariant, consistent with GSPO). For KL metrics, aggregation is arithmetic: 'sum' = sum(kl_i), 'mean' = mean(kl_i). 'max': max of per-token metric values (most conservative). **Choices:** `sum`, `mean`, `max` | +| `upper` | float | `5.0` | Upper bound for filtering. Tokens/sequences with metric > upper are filtered out (loss_mask zeroed). For 'ratio' metric: must be > 1.0, typical values are 2.0 or 5.0. For 'kl_k2'/'kl_k3' metrics: typical values are 0.5-2.0. | +| `lower` | float \| None | `None` | Lower bound for filtering (optional). None means no lower bound. For 'ratio' metric: typical value is 0.5 (filter out tokens where policy probability dropped significantly). Must be > 0. For 'kl_k1' metric: can be used to filter negative KL estimates. | (section-scheduler)= diff --git a/docs/en/reference/rollout_workflow.md b/docs/en/reference/rollout_workflow.md index b1ed9549d6..9b0720d595 100644 --- a/docs/en/reference/rollout_workflow.md +++ b/docs/en/reference/rollout_workflow.md @@ -140,16 +140,37 @@ Each line in the JSONL file contains: { "task_id": 42, "sample_idx": 0, - "seqlen": 256, + "seqlen": 512, "prompt_len": 128, "head_version": 5, - "tail_version": 5, + "tail_version": 6, + "version_rle": [[5, 100], [6, 200]], "reward": 1.0, "prompt": "<|im_start|>user\nWhat is 2+2?<|im_end|>\n<|im_start|>assistant\n", "completion": "The answer is 4.<|im_end|>" } ``` +**Field descriptions:** + +| Field | Description | +| -------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `task_id` | Batch task identifier | +| `sample_idx` | Index of the sample within the batch | +| `seqlen` | Effective sequence length | +| `prompt_len` | Index of the first generated token (`mask.index(1)`). For multi-turn agent rollouts this is the position of the first assistant generation, not `seqlen - sum(mask)` | +| `head_version` | Per-sample minimum model version among `loss_mask==1` tokens | +| `tail_version` | Per-sample maximum model version among `loss_mask==1` tokens | +| `version_rle` | Run-length encoded per-token version sequence (output tokens only), e.g. `[[5, 100], [6, 200]]` | +| `reward` | Reward value for this sample | +| `prompt` | Decoded prompt text | +| `completion` | Decoded completion text | +| `segments` | *(multi-turn only)* List of `{"role": "prompt"\|"gen"\|"context", "len": N, "text": "..."}` | + +**Directory naming:** The `{version}` directory is named by the batch-global maximum +version (`global_tail`). Individual records within the same directory may have +`tail_version <= global_tail`. + ## Grouped Rollout Grouped rollout runs the same workflow multiple times per input prompt, producing diff --git a/docs/en/tutorial/installation.md b/docs/en/tutorial/installation.md index 8e2a4de9ca..cd6a846023 100644 --- a/docs/en/tutorial/installation.md +++ b/docs/en/tutorial/installation.md @@ -24,7 +24,7 @@ The following hardware configuration has been extensively tested: | Git LFS | Required for downloading models, datasets, and AReaL code. See [installation guide](https://docs.github.com/en/repositories/working-with-files/managing-large-files/installing-git-large-file-storage) | | Docker | 27.5.1 | | NVIDIA Container Toolkit | See [installation guide](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) | -| AReaL Image | `ghcr.io/areal-project/areal-runtime:v1.0.4-sglang` (default) or `v1.0.4-vllm`. Includes runtime dependencies and Ray components. | +| AReaL Image | `ghcr.io/areal-project/areal-runtime:v2.0.0-sglang` (default) or `v2.0.0-vllm`. Includes runtime dependencies and Ray components. | **Note**: This tutorial does not cover the installation of NVIDIA Drivers, CUDA, or shared storage mounting, as these depend on your specific node configuration and system @@ -42,11 +42,11 @@ We recommend using Docker with our provided image. The Dockerfile is available i top-level directory of the AReaL repository. ```bash -docker pull ghcr.io/areal-project/areal-runtime:v1.0.4-sglang +docker pull ghcr.io/areal-project/areal-runtime:v2.0.0-sglang docker run -it --name areal-node1 \ --privileged --gpus all --network host \ --shm-size 700g -v /path/to/mount:/path/to/mount \ - ghcr.io/areal-project/areal-runtime:v1.0.4-sglang \ + ghcr.io/areal-project/areal-runtime:v2.0.0-sglang \ /bin/bash git clone https://github.com/areal-project/AReaL /path/to/mount/AReaL cd /path/to/mount/AReaL @@ -54,7 +54,7 @@ uv pip install -e . --no-deps ``` A vLLM variant of the Docker image is also available at -`ghcr.io/areal-project/areal-runtime:v1.0.4-vllm`. Replace the image tag in the commands +`ghcr.io/areal-project/areal-runtime:v2.0.0-vllm`. Replace the image tag in the commands above if you prefer vLLM as the inference backend. ### Option 2: Custom Environment Installation diff --git a/docs/paper/AReaL2.0_report.pdf b/docs/paper/AReaL2.0_report.pdf new file mode 100644 index 0000000000..cb720dbf88 Binary files /dev/null and b/docs/paper/AReaL2.0_report.pdf differ diff --git a/docs/zh/algorithms/grpo_series.md b/docs/zh/algorithms/grpo_series.md index d882b8a6e8..a9cf12f5aa 100644 --- a/docs/zh/algorithms/grpo_series.md +++ b/docs/zh/algorithms/grpo_series.md @@ -17,6 +17,10 @@ - **SAPO**: [论文](https://arxiv.org/abs/2511.20347) - **GSPO** (Qwen3): [论文](https://arxiv.org/abs/2507.18071)， [博客](https://qwenlm.github.io/blog/gspo/) +- **IcePop**：[博客](https://ringtech.notion.site/icepop) — 基于重要性比率的token掩码（可与其它RL算法组合使用） +- **KPop**：[博客](https://ringtech.notion.site/kpop) — 双向二元KL散度token掩码（可与其它RL算法组合使用） + +IcePop和KPop是token掩码策略，可以与上面列出的任意RL算法组合使用。 这些算法共享相同的基础目标，但在归一化策略、裁剪机制、重要性采样级别等方面有所不同。通过调整AReaL中的少量配置参数，你可以在不同算法之间切换。 @@ -30,7 +34,7 @@ | **ray** | `python3 examples/math/gsm8k_rl.py --config examples/math/gsm8k_.yaml scheduler.type=ray` | | **slurm** | `python3 examples/math/gsm8k_rl.py --config examples/math/gsm8k_.yaml scheduler.type=slurm` | -将 `` 替换为：`ppo`、`grpo`、`drgrpo`、`liteppo`、`rloo`、`gspo`、`dapo_dynamic_bs` 或 `sapo`。 +将 `` 替换为：`ppo`、`grpo`、`drgrpo`、`liteppo`、`rloo`、`gspo`、`dapo_dynamic_bs`、`sapo`、`icepop` 或 `kpop`。 ### 通过CLI覆盖切换算法 @@ -134,8 +138,10 @@ actor: | **LitePPO** | `group` | `batch` | `false` | `token` | - | | **RLOO** | `group` | `null` | `true` | `token` | - | | **GSPO** | `batch` | `batch` | `false` | `sequence` | - | -| **DAPO** | `batch` | `batch` | `false` | `token` | 非对称裁剪，动态采样 | -| **SAPO** | `batch` | `batch` | `false` | `token` | `use_sapo_loss=true` | +| **DAPO** | `batch` | `batch` | `false` | `token` | 非对称裁剪，动态采样 | +| **SAPO** | `batch` | `batch` | `false` | `token` | `use_sapo_loss=true` | +| **IcePop** | `batch` | `batch` | `false` | `token` | `rejection_sampling.metric=ratio` | +| **KPop** | `batch` | `batch` | `false` | `token` | `rejection_sampling.metric=binary_kl` | **注意**："GRPO"行反映原始DeepSeekMath公式。AReaL的默认GRPO配置使用这些设置，但已移除长度归一化（见下文AReaL实现说明）。 @@ -271,6 +277,62 @@ trainer.train( | ------------ | ---- | ------- | ---------------- | | `dynamic_bs` | bool | `false` | 启用动态批量大小 | +### IcePop + +IcePop对重要性比率 $r_{i,t} = \frac{\pi_\theta(o_{i,t} \mid q, o_{i, \phi$ 的token被掩码（此处 $\phi$ 对应 `actor.rejection_sampling.upper`）。 + +通过 `rejection_sampling` 配置的 `metric=binary_kl` 实现： + +```yaml +actor: + use_decoupled_loss: true + rejection_sampling: + level: token + action: mask + metric: binary_kl + upper: 2.0 +``` + +| 参数 | 类型 | 默认值 | 描述 | +| ------------------------------------ | ----------- | ------ | ------------------------------ | +| `actor.rejection_sampling.metric` | str | - | 设置为 `binary_kl` 以启用KPop | +| `actor.rejection_sampling.upper` | float | `2.0` | KL散度阈值（$\phi$） | + +**注意：** KPop仅支持 `action=mask`（不支持 `clamp`），且 `lower` 在 `binary_kl` 中不使用。KPop需要 `actor.use_decoupled_loss=true`，否则 `rejection_sampling` 不生效。 + +完整的配置示例见 `examples/math/gsm8k_kpop.yaml`。 + ## 核心概念 **奖励**：AReaL假设基于结果的奖励。每个可能由连接的LLM输入-输出对组成的轨迹，在序列级而非token级被分配一个标量奖励。 diff --git a/docs/zh/cli_reference.md b/docs/zh/cli_reference.md index d8cb5c8f80..53a04daa75 100644 --- a/docs/zh/cli_reference.md +++ b/docs/zh/cli_reference.md @@ -117,6 +117,7 @@ Base configuration class for all experiment types with common settings. | `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | | `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | | `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | +| `post_exit_hook` | string | `""` | Shell command run after launcher shutdown. LOG_DIR is injected; failures are logged and ignored. | (section-grpo)= @@ -147,6 +148,7 @@ A dummy place holder of GRPO config for backward compatibility. | `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | | `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | | `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | +| `post_exit_hook` | string | `""` | Shell command run after launcher shutdown. LOG_DIR is injected; failures are logged and ignored. | | `gconfig` | [`GenerationHyperparameters`](section-generation-hyperparameters) | **Required** | - | | `eval_gconfig` | [`GenerationHyperparameters`](section-generation-hyperparameters) \| None | `None` | Generation hyperparameters for evaluation. If None, use gconfig. | | `rollout` | [`InferenceEngineConfig`](section-inference-engine) | **Required** | - | @@ -185,6 +187,7 @@ Configuration for Proximal Policy Optimization (PPO) reinforcement learning expe | `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | | `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | | `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | +| `post_exit_hook` | string | `""` | Shell command run after launcher shutdown. LOG_DIR is injected; failures are logged and ignored. | | `gconfig` | [`GenerationHyperparameters`](section-generation-hyperparameters) | **Required** | - | | `eval_gconfig` | [`GenerationHyperparameters`](section-generation-hyperparameters) \| None | `None` | Generation hyperparameters for evaluation. If None, use gconfig. | | `rollout` | [`InferenceEngineConfig`](section-inference-engine) | **Required** | - | @@ -223,6 +226,7 @@ Configuration for Reward Model (RW) training experiments. | `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | | `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | | `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | +| `post_exit_hook` | string | `""` | Shell command run after launcher shutdown. LOG_DIR is injected; failures are logged and ignored. | | `actor` | [`TrainEngineConfig`](section-train-engine) | **Required** | - | (section-sft)= @@ -254,6 +258,7 @@ Configuration for Supervised Fine-Tuning (SFT) experiments. | `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | | `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | | `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | +| `post_exit_hook` | string | `""` | Shell command run after launcher shutdown. LOG_DIR is injected; failures are logged and ignored. | | `actor` | [`TrainEngineConfig`](section-train-engine) | **Required** | - | (section-fsdp-engine)= @@ -399,11 +404,12 @@ Configuration for PPO actor model, a subclass of a TrainEngine. | `use_sapo_loss` | boolean | `False` | Use SAPO loss (mutually exclusive with PPO clipping) | | `sapo_tau_pos` | float | `1.0` | SAPO temperature for positive advantages | | `sapo_tau_neg` | float | `1.05` | SAPO temperature for negative advantages | +| `use_cispo_loss` | boolean | `False` | Use CISPO loss: clip the importance-sampling weight under stop-gradient and keep gradient on every token's log pi (MiniMax-M1 Eq. 4-5). Mutually exclusive with SAPO. Token-level only. Requires eps_clip_higher > 0; recommended eps_clip=1.0 (single-sided, lower bound 0) with eps_clip_higher=4.0. | | `recompute_logprob` | boolean | `False` | Recompute log probability and replace the log probability returned by inference. | | `use_decoupled_loss` | boolean | `False` | Use the decoupled loss. Implicitly enables recompute_logprob. | | `rejection_sampling` | [`RejectionSamplingConfig`](section-rejection-sampling) \| None | `None` | Rejection sampling configuration for filtering stale samples. None disables filtering (equivalent to old behave_imp_weight_mode='disabled'). Only effective when use_decoupled_loss=True. | | `importance_sampling_level` | string | `"token"` | Level at which to compute importance sampling ratios. 'token': per-token ratios (standard PPO). 'sequence': sequence-level geometric mean of per-token ratios (GSPO). **Choices:** `token`, `sequence` | -| `prox_logp_method` | string | `"recompute"` | Method for computing proximal policy log-probabilities in decoupled PPO. Only effective when use_decoupled_loss=True. Options: 'recompute' (default): Standard decoupled PPO, recompute proximal policy via forward pass. 'loglinear': Use log-linear interpolation to approximate proximal policy (skip forward pass). 'metrics': Like 'recompute', but also compute approximation metrics for evaluation. **Choices:** `recompute`, `loglinear`, `metrics` | +| `prox_logp_method` | string | `"recompute"` | Method for computing proximal policy log-probabilities in decoupled PPO. Only effective when use_decoupled_loss=True. Options: 'recompute' (default): Standard decoupled PPO, recompute proximal policy via forward pass. 'loglinear': Use log-linear interpolation to approximate proximal policy (skip forward pass). 'metrics': Like 'recompute', but also compute approximation metrics for evaluation. 'reuse_train_logp': Reuse training forward-pass logprobs as the proximal logp (skip the extra forward; requires ppo_n_minibatches=1). **Choices:** `recompute`, `loglinear`, `metrics`, `reuse_train_logp` | | `log_agent_stats` | boolean | `False` | Log statistics for agent trajectories | | `log_agent_stats_keys` | list of string | **Required** | Keys for logging agent trajectory statistics | | `max_new_tokens` | integer | `1024` | Maximum number of new tokens to generate | @@ -891,6 +897,8 @@ Consolidates proxy settings (mode, parsers, export) with agent-service orchestra | `engine_max_tokens` | integer \| None | `None` | Maximum total tokens for the engine (prompt + completion). | | `turn_discount` | float | `1.0` | Discount factor for multi-turn reward propagation. | | `export_style` | string | `"individual"` | Export style: 'individual' (all interactions) or 'concat' (leaf nodes only). The 'individual' style exports each interaction (input-output-reward) step separately, and treats them as independent samples to train the model. The 'concat' style exports only the final concatenated trajectory from the root. It is only suitable for linear conversation histories without token mismatching (whether valid depends on the tokenizer). **Choices:** `individual`, `concat` | +| `message_preprocessors` | list of string | **Required** | List of message preprocessor class paths applied, in order, to Anthropic-compatible `/v1/messages` requests after translating them to OpenAI-compatible requests. Native OpenAI `/chat/completions` and `/responses` requests are not preprocessed. Each entry is a dotted import path to a callable class. | +| `prefix_matcher` | string \| None | `None` | Dotted import path to a custom prefix matcher function for InteractionCache parent-child matching. The function must accept two list\[dict\] arguments (candidate prefix, full messages) and return bool. When None, exact element-wise equality is used. | | `subproc_max_workers` | integer | `4` | Maximum number of worker processes for subprocess mode execution pool. | | `session_timeout_seconds` | integer | `3600` | Session timeout in seconds. Sessions inactive longer than this will be garbage collected. | | `set_reward_finish_timeout` | float | `0.0` | Timeout in seconds to wait for additional reward updates before finalizing a session. | @@ -962,6 +970,7 @@ Configuration for Direct Preference Optimization (DPO) experiments. | `sglang` | [`SGLangConfig`](section-sg-lang) | **Required** | - | | `vllm` | [`vLLMConfig`](section-v-llm) | **Required** | - | | `scheduler` | [`SchedulerConfig`](section-scheduler) | **Required** | - | +| `post_exit_hook` | string | `""` | Shell command run after launcher shutdown. LOG_DIR is injected; failures are logged and ignored. | | `actor` | [`DPOEngineConfig`](section-dpo-engine) | **Required** | - | | `ref` | [`DPOEngineConfig`](section-dpo-engine) | **Required** | - | @@ -1094,10 +1103,17 @@ Refer to Megatron-LM documentation for implementation details. | `moe_enable_deepep` | boolean | `False` | - | | `moe_token_dispatcher_type` | string | `"alltoall"` | Type of token dispatcher. Options: 'allgather','alltoall' and 'flex'. | | `moe_permute_fusion` | boolean | `False` | Fuse token rearrangement ops during token dispatching. | +| `moe_router_fusion` | boolean | `False` | Enable fusion for MoE TopK routing and aux-loss computation. Requires TransformerEngine >= 2.7.0. | +| `moe_router_bias_update_rate` | float | `0.0` | Update rate for auxiliary-loss-free MoE load balancing (DeepSeek V3 style). Controls how fast expert_bias adjusts. Default 0.0 disables bias updates; set a positive value such as 1e-3 to enable. | +| `moe_z_loss_coeff` | float \| None | `None` | Scaling coefficient for router z-loss. Complements auxiliary-loss-free load balancing for router stability. A starting value of 1e-3 is recommended. None disables z-loss. | +| `enable_fp32_lm_head` | boolean | `False` | Cast lm_head output to FP32 before loss computation for numerical stability. | +| `cross_entropy_loss_fusion` | boolean | `False` | Enable fused cross-entropy loss kernel for better performance. | | `fp8_config` | [`FP8EngineConfig`](section-fp8-engine) \| None | `None` | - | | `bridge_type` | string | `"mbridge"` | Bridge backend for MegatronEngine. Choices: 'mbridge' or 'megatron-bridge'. **Choices:** `mbridge`, `megatron-bridge` | | `use_mbridge_save` | boolean | `False` | Use mbridge's save method to save gpu memory when saving weights. | | `use_bridge_for_update_weights` | boolean | `False` | When True and bridge_type='megatron-bridge', delegate live weight sync to bridge.export_hf_weights instead of the hand-rolled convert_to_hf registry. Required for models without a registry entry (e.g. Qwen3.5). FP8 paths fall back to the registry automatically. | +| `disable_grad_buffers_cpu_backup` | boolean | `False` | When offloading with torch_memory_saver, skip CPU backup for Megatron gradient buffers (they are recomputed each step). | +| `enable_mtp` | boolean | `False` | Keep the model's Multi-Token-Prediction (MTP) head (bridge_type=megatron-bridge only). Default False drops it. | (section-memory-profiler)= @@ -1161,14 +1177,14 @@ Attributes: lower: Lower bound for filtering (optional). ``` -| Parameter | Type | Default | Description | -| --------- | ------------- | --------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `level` | string | `"token"` | Filtering granularity. 'token': per-token filtering (each token judged independently). 'sequence': per-sequence filtering (all tokens in a sequence share the same fate). When metric='ratio', both the filtering decision and the correction weight (behave_imp_weight) operate at sequence level using the geometric mean. **Choices:** `token`, `sequence` | -| `action` | string | `"mask"` | Action to take when metric exceeds threshold. 'mask': zero out loss_mask for filtered tokens/sequences (rejection, completely excludes from gradient computation). 'clamp': clamp importance weight to \[lower, upper\] bounds (truncation, tokens still participate in gradient but with bounded weight). **Choices:** `mask`, `clamp` | -| `metric` | string | `"ratio"` | Divergence metric for filtering. 'ratio': direct importance ratio π_proximal/π_behave. 'kl_k1': KL estimator k1 = log(r), forward KL unbiased estimator (can be negative). 'kl_k2': KL estimator k2 = 0.5 * (log r)^2, non-negative quadratic approximation. 'kl_k3': KL estimator k3 = r - log(r) - 1, non-negative exact forward KL estimator. **Choices:** `ratio`, `kl_k1`, `kl_k2`, `kl_k3` | -| `agg` | string | `"mean"` | Aggregation method for sequence-level filtering. Only used when level='sequence'. For 'ratio' metric, aggregation is in log space: 'sum' = exp(sum(log(r_i))), 'mean' = exp(mean(log(r_i))) = geometric mean (length-invariant, consistent with GSPO). For KL metrics, aggregation is arithmetic: 'sum' = sum(kl_i), 'mean' = mean(kl_i). 'max': max of per-token metric values (most conservative). **Choices:** `sum`, `mean`, `max` | -| `upper` | float | `5.0` | Upper bound for filtering. Tokens/sequences with metric > upper are filtered out (loss_mask zeroed). For 'ratio' metric: must be > 1.0, typical values are 2.0 or 5.0. For 'kl_k2'/'kl_k3' metrics: typical values are 0.5-2.0. | -| `lower` | float \| None | `None` | Lower bound for filtering (optional). None means no lower bound. For 'ratio' metric: typical value is 0.5 (filter out tokens where policy probability dropped significantly). Must be > 0. For 'kl_k1' metric: can be used to filter negative KL estimates. | +| Parameter | Type | Default | Description | +| --------- | ------------- | --------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `level` | string | `"token"` | Filtering granularity. 'token': per-token filtering (each token judged independently). 'sequence': per-sequence filtering (all tokens in a sequence share the same fate). When metric='ratio', both the filtering decision and the correction weight (behave_imp_weight) operate at sequence level using the geometric mean. **Choices:** `token`, `sequence` | +| `action` | string | `"mask"` | Action to take when metric exceeds threshold. 'mask': zero out loss_mask for filtered tokens/sequences (rejection, completely excludes from gradient computation). 'clamp': clamp importance weight to \[lower, upper\] bounds (truncation, tokens still participate in gradient but with bounded weight). **Choices:** `mask`, `clamp` | +| `metric` | string | `"ratio"` | Divergence metric for filtering. 'ratio': direct importance ratio π_proximal/π_behave. 'kl_k1': KL estimator k1 = log(r), forward KL unbiased estimator (can be negative). 'kl_k2': KL estimator k2 = 0.5 * (log r)^2, non-negative quadratic approximation. 'kl_k3': KL estimator k3 = r - log(r) - 1, non-negative exact forward KL estimator. 'binary_kl': KPop (symmetric binary KL divergence) — masks tokens where either KL(proximal | +| `agg` | string | `"mean"` | Aggregation method for sequence-level filtering. Only used when level='sequence'. For 'ratio' metric, aggregation is in log space: 'sum' = exp(sum(log(r_i))), 'mean' = exp(mean(log(r_i))) = geometric mean (length-invariant, consistent with GSPO). For KL metrics, aggregation is arithmetic: 'sum' = sum(kl_i), 'mean' = mean(kl_i). 'max': max of per-token metric values (most conservative). **Choices:** `sum`, `mean`, `max` | +| `upper` | float | `5.0` | Upper bound for filtering. Tokens/sequences with metric > upper are filtered out (loss_mask zeroed). For 'ratio' metric: must be > 1.0, typical values are 2.0 or 5.0. For 'kl_k2'/'kl_k3' metrics: typical values are 0.5-2.0. | +| `lower` | float \| None | `None` | Lower bound for filtering (optional). None means no lower bound. For 'ratio' metric: typical value is 0.5 (filter out tokens where policy probability dropped significantly). Must be > 0. For 'kl_k1' metric: can be used to filter negative KL estimates. | (section-scheduler)= diff --git a/docs/zh/reference/rollout_workflow.md b/docs/zh/reference/rollout_workflow.md index f999b31875..da3a791189 100644 --- a/docs/zh/reference/rollout_workflow.md +++ b/docs/zh/reference/rollout_workflow.md @@ -134,16 +134,36 @@ JSONL 文件的每一行包含： { "task_id": 42, "sample_idx": 0, - "seqlen": 256, + "seqlen": 512, "prompt_len": 128, "head_version": 5, - "tail_version": 5, + "tail_version": 6, + "version_rle": [[5, 100], [6, 200]], "reward": 1.0, "prompt": "<|im_start|>user\nWhat is 2+2?<|im_end|>\n<|im_start|>assistant\n", "completion": "The answer is 4.<|im_end|>" } ``` +**字段说明：** + +| 字段 | 说明 | +| -------------- | --------------------------------------------------------------------------------------------------------------------- | +| `task_id` | Batch 任务标识符 | +| `sample_idx` | Batch 内的样本索引 | +| `seqlen` | 有效序列长度 | +| `prompt_len` | 第一个生成 token 的位置（`mask.index(1)`）。多轮 agent rollout 中为首次 assistant 生成位置，而非 `seqlen - sum(mask)` | +| `head_version` | Per-sample 最小模型版本，仅统计 `loss_mask==1` 的 token | +| `tail_version` | Per-sample 最大模型版本，仅统计 `loss_mask==1` 的 token | +| `version_rle` | RLE 编码的 per-token 版本序列（仅输出 token），例如 `[[5, 100], [6, 200]]` | +| `reward` | 该样本的 reward 值 | +| `prompt` | 解码后的 prompt 文本 | +| `completion` | 解码后的 completion 文本 | +| `segments` | *（仅多轮）* 分段列表 `[{"role": "prompt"\|"gen"\|"context", "len": N, "text": "..."}]` | + +**目录命名：** `{version}` 目录以 batch 内的全局最大版本（`global_tail`）命名。 同一目录内的单条 record 的 +`tail_version` 可能 \<= 目录名。 + ## 分组 Rollout 分组 Rollout 对每个输入提示多次运行相同工作流，生成多样化回复用于训练。这对于像 GRPO 这样受益于每个提示多个样本的算法很有用。 diff --git a/docs/zh/tutorial/installation.md b/docs/zh/tutorial/installation.md index ca6ce017e1..3027291164 100644 --- a/docs/zh/tutorial/installation.md +++ b/docs/zh/tutorial/installation.md @@ -24,7 +24,7 @@ | Git LFS | 用于下载模型、数据集和 AReaL 代码。请参阅[安装指南](https://docs.github.com/en/repositories/working-with-files/managing-large-files/installing-git-large-file-storage) | | Docker | 27.5.1 | | NVIDIA Container Toolkit | 请参阅[安装指南](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) | -| AReaL 镜像 | `ghcr.io/areal-project/areal-runtime:v1.0.4-sglang`（默认）或 `v1.0.4-vllm`。包含运行时依赖和 Ray 组件。 | +| AReaL 镜像 | `ghcr.io/areal-project/areal-runtime:v2.0.0-sglang`（默认）或 `v2.0.0-vllm`。包含运行时依赖和 Ray 组件。 | **注意**：本教程不涵盖 NVIDIA 驱动、CUDA 或共享存储挂载的安装，因为这些取决于您具体的节点配置和系统版本。请独立完成这些安装。 @@ -37,18 +37,18 @@ 我们推荐使用 Docker 和提供的镜像。Dockerfile 位于 AReaL 仓库的顶级目录。 ```bash -docker pull ghcr.io/areal-project/areal-runtime:v1.0.4-sglang +docker pull ghcr.io/areal-project/areal-runtime:v2.0.0-sglang docker run -it --name areal-node1 \ --privileged --gpus all --network host \ --shm-size 700g -v /path/to/mount:/path/to/mount \ - ghcr.io/areal-project/areal-runtime:v1.0.4-sglang \ + ghcr.io/areal-project/areal-runtime:v2.0.0-sglang \ /bin/bash git clone https://github.com/areal-project/AReaL /path/to/mount/AReaL cd /path/to/mount/AReaL uv pip install -e . --no-deps ``` -vLLM 变体的 Docker 镜像也可使用： `ghcr.io/areal-project/areal-runtime:v1.0.4-vllm`。如果您偏好使用 vLLM +vLLM 变体的 Docker 镜像也可使用： `ghcr.io/areal-project/areal-runtime:v2.0.0-vllm`。如果您偏好使用 vLLM 作为推理后端，请将上述命令中的镜像标签替换为该变体。 ### 方式 2：自定义环境安装 diff --git a/examples/agent_service/README.md b/examples/agent_service/README.md index 33e1deca5d..73371cfb45 100644 --- a/examples/agent_service/README.md +++ b/examples/agent_service/README.md @@ -84,7 +84,7 @@ Turn 2: Client → Gateway → Router (same DataProxy) → DataProxy → Worker import os from areal.api.cli_args import AgentConfig, SchedulingSpec -from areal.experimental.agent_service.controller import ( +from areal.v2.agent_service.controller import ( AgentController, ) from areal.infra.scheduler.local import LocalScheduler diff --git a/examples/agent_service/agent.py b/examples/agent_service/agent.py index c05f3bebe5..d6287038d6 100644 --- a/examples/agent_service/agent.py +++ b/examples/agent_service/agent.py @@ -30,12 +30,12 @@ ToolUseBlock, ) -from areal.experimental.agent_service.types import ( +from areal.utils import logging +from areal.v2.agent_service.types import ( AgentRequest, AgentResponse, EventEmitter, ) -from areal.utils import logging logger = logging.getLogger("ClaudeAgent") diff --git a/examples/agent_service/run_agent_service.py b/examples/agent_service/run_agent_service.py index 8f97ccbf32..799b40e4a0 100644 --- a/examples/agent_service/run_agent_service.py +++ b/examples/agent_service/run_agent_service.py @@ -22,7 +22,7 @@ import httpx from areal.api.cli_args import AgentConfig -from areal.experimental.agent_service.controller import AgentController +from areal.v2.agent_service.controller import AgentController async def _wait_healthy(url: str, timeout: float = 60.0) -> None: diff --git a/examples/experimental/inference_service/online_rollout.py b/examples/experimental/inference_service/online_rollout.py index a2f2d81677..74ee29c3ea 100644 --- a/examples/experimental/inference_service/online_rollout.py +++ b/examples/experimental/inference_service/online_rollout.py @@ -21,11 +21,11 @@ def main(args: list[str]) -> None: ext_args, remaining = parser.parse_known_args(args) from areal.api.cli_args import PPOConfig, load_expr_config - from areal.experimental.inference_service.controller.controller import ( - RolloutControllerV2, - ) from areal.utils import logging from areal.utils.environ import is_single_controller + from areal.v2.inference_service.controller.controller import ( + RolloutControllerV2, + ) logger = logging.getLogger("InferenceServiceOnlineTrain") diff --git a/examples/experimental/inference_service/tau2_rollout.py b/examples/experimental/inference_service/tau2_rollout.py index b78701b023..7bdfec0ef6 100644 --- a/examples/experimental/inference_service/tau2_rollout.py +++ b/examples/experimental/inference_service/tau2_rollout.py @@ -28,10 +28,10 @@ TrainDatasetConfig, load_expr_config, ) -from areal.experimental.inference_service.controller.controller import ( +from areal.utils import logging +from areal.v2.inference_service.controller.controller import ( RolloutControllerV2, ) -from areal.utils import logging logger = logging.getLogger("Tau2GatewayRollout") diff --git a/examples/experimental/inference_service/tau2_rollout.yaml b/examples/experimental/inference_service/tau2_rollout.yaml index d76bab14aa..c467c56a48 100644 --- a/examples/experimental/inference_service/tau2_rollout.yaml +++ b/examples/experimental/inference_service/tau2_rollout.yaml @@ -56,7 +56,7 @@ rollout: gpu: 1 cpu: 2 mem: 16 - cmd: python3 -m areal.experimental.inference_service.guard + cmd: python3 -m areal.v2.inference_service.guard env_vars: AREAL_PROXY_WARN_ONCE: "1" diff --git a/examples/hermes/README.md b/examples/hermes/README.md new file mode 100644 index 0000000000..7c01ef28a3 --- /dev/null +++ b/examples/hermes/README.md @@ -0,0 +1,192 @@ +# Agent Service — Hermes + +## Overview + +This example runs the **Hermes agent** +([Nous Research `hermes-agent`](https://github.com/nousresearch/hermes-agent)) inside +AReaL's Agent Service. Hermes ships as a Python library, so the Worker instantiates +**one in-process Hermes `AIAgent` per session** directly inside the Worker process — +**you never launch Hermes yourself**. The Agent Service is started with the +`areal agent run` CLI; you then interact and (optionally) train through a small set of +single-purpose scripts. + +``` +Client → Gateway (HTTP) → Router → DataProxy (session state) → Worker (Hermes AIAgent) +``` + +Each user message becomes one turn of the Hermes conversation; the per-session `AIAgent` +drives its configured OpenAI-compatible upstream LLM internally. AReaL's DataProxy owns +the conversation history and replays it into every turn, so the `AIAgent` is kept +**stateless across turns** (`skip_memory=True`, `skip_context_files=True`, +`session_db=None`). A consequence of this design: Hermes' own `memory` tool reports +"Memory is not available" — that is expected, since cross-turn context comes from the +DataProxy's replayed history, not from Hermes' persistence. + +This directory also contains the **RL training flow** (`train.py` + `config.yaml`). For +training, run `train.py` to bring up AReaL's trainer plus an inference gateway, then +point the Hermes turns at that gateway (via `hermes_loop.py`'s inference-routing flags) +so every LLM call is captured as a training trajectory. + +> **See also** +> +> - [Agentic RL tutorial](../../docs/en/tutorial/agentic_rl.md) — background on how +> AReaL trains agents +> - [Custom agent workflows](../../docs/en/customization/agent.md) — how to integrate +> your own agent framework +> - [Agent workflow reference](../../docs/en/reference/agent_workflow.md) — internal +> architecture details + +**Disclaimer**: RL-finetuned models may exhibit unexpected behaviors. Please ensure +strict permission rules and an isolated execution environment for your agent runtime. + +## How it fits together + +``` +┌──────────────────────────────────────┐ LLM calls (self-evolution) ┌────────────────────────┐ +│ Agent Service │ ─────────────────────────────▶ │ AReaL inference gateway│ +│ (areal agent run) │ inf_base_url = http:// │ (started by train.py) │ +│ Gateway/Router/DataProxy/Worker │ session key = sk-sess-* │ │ +│ + in-process Hermes AIAgent │ ◀───────────────────────────── │ records tokens + │ +└──────────────────────────────────────┘ model output │ logprobs → RL │ + ▲ └────────────────────────┘ + │ hermes_loop.py │ + │ (the interactive "You:" prompt) set_reward.py (score the trajectory) +``` + +One **episode** = the turns collected under a single per-session `sk-sess-*` key (minted +by `start_session.py`). You score it with `set_reward.py`, then start the next episode. + +## Prerequisites + +### 1. GPUs (for RL training only) + +A GPU machine with at least **2 NVIDIA GPUs** (compute capability 8.0 or higher, i.e. +Ampere / Hopper). Not required if you only run the agent against an env upstream LLM +(plain chat, no training). + +### 2. Install Hermes into AReaL's venv + +Hermes' top-level module is `run_agent`. The Worker process is forked with +`sys.executable` (the interpreter you launch the controller with), so **`areal` and +`run_agent` must be importable from the same venv**: + +```bash +uv pip install hermes-agent +python -c "import areal; from run_agent import AIAgent; print('co-import OK')" +``` + +A bare `hermes-agent` install is moderate, **not** heavy: every large optional +integration (`anthropic`, `slack`, `matrix`, `modal`, browser/messaging, …) sits behind +a `pip install hermes-agent[extra]` marker and is not pulled in. No torch/CUDA/ML +packages are added. + +> **Gotchas** +> +> - **Run with the project `.venv` python directly** (`python` / `.venv/bin/python`), +> **not `uv run`** — `uv run` re-syncs the env to `uv.lock` and resets the shared +> packages (`openai`, `pydantic`, `rich`, …) to AReaL's pinned versions. +> - **Do not run `uv sync`** while you need Hermes — it removes `hermes-agent` (it is +> not in `uv.lock`). Re-add with `uv pip install hermes-agent`. + +## Quick start — plain chat + +The five steps below are run from the repo root. + +### Step 1 — Start the training service (embeds the inference gateway) + +`config.yaml` holds the defaults (v2 controllers, 1 node × 2 GPUs, `batch_size=1`, admin +keys); CLI flags override it. The explicit form below just documents those defaults: + +```bash +uv run python3 examples/hermes/train.py \ + --config examples/hermes/config.yaml \ + actor.path=/path/to/your_model \ + actor.admin_api_key=sk-123456 \ + rollout.admin_api_key=sk-123456 +``` + +Note this address from the logs — it is your `` below: + +``` +Proxy gateway available at http://X.X.X.X:PORT +``` + +> **Key wiring** +> +> `rollout.admin_api_key` is the **inference gateway** admin key — reuse it for +> `start_session.py --admin-key` and `set_reward.py` (Steps 3 and 5). +> `actor.admin_api_key` is the trainer/actor key, unused by the interaction scripts. See +> the [CLI reference](../../docs/en/cli_reference.md) and +> [allocation mode reference](../../docs/en/reference/alloc_mode.md) for fields and GPU +> layout. + +### Step 2 — Start the Hermes Agent Service + +Same command as the quick start. The agent's env upstream is optional here — once +self-evolution fields are supplied per turn (Step 4), the inference gateway upstream +takes over. + +```bash +areal agent run \ + --service default \ + --agent examples.hermes.hermes.HermesAgent \ + --num-pairs 1 \ + --admin-api-key sk-123456 +``` + +Note the printed `` address. + +### Step 3 — Start a session on the inference gateway + +Copy the printed `sk-sess-*` key — forward it to the agent (Step 4) and score the +episode with it (Step 5). To reuse the key for the next episode (auto-ends and exports +the previous one): + +```bash +python examples/hermes/start_session.py http:// --admin-key sk-123456 +``` + +These are **your own upstream LLM credentials** (the agent's fallback chat backend), +**not** the `sk-sess-*` key returned above — fill in your provider's values: + +```bash +export HERMES_UPSTREAM_BASE_URL="https://your-llm/v1" +export HERMES_UPSTREAM_API_KEY="your-upstream-api-key" +export HERMES_UPSTREAM_MODEL="your-model" +``` + +### Step 4 — Interact (produces a trajectory) + +Forward the inference-routing flags so the agent's LLM calls flow through the inference +gateway under your session key and get captured. You **must** actually interact, or the +episode has no data. `` is the `sk-sess-*` key returned by +`start_session.py` in Step 3. + +```bash +python examples/hermes/hermes_loop.py http:// \ + --admin-api-key sk-123456 \ + --inf-base-url http:// \ + --session-api-key +``` + +### Step 5 — Score the episode + +Use the same `sk-sess-*` key from Step 3 as `--api-key`: + +```bash +python examples/hermes/set_reward.py http:// \ + --api-key --reward 1.0 +``` + +Keep the reward in **\[-1, 1\]** for training stability. + +## Files + +| File | Description | +| ------------------ | ---------------------------------------------------------------- | +| `hermes.py` | `HermesAgent` — in-process per-session Hermes `AIAgent` runnable | +| `hermes_loop.py` | Standalone interactive `You:` prompt against the agent gateway | +| `start_session.py` | Mint a per-session `sk-sess-*` key on the inference gateway | +| `set_reward.py` | Assign a scalar reward to a session's trajectory | +| `train.py` | RL trainer entry point (embeds the inference gateway) | +| `config.yaml` | Training configuration (v2 controllers, 2-GPU defaults) | diff --git a/examples/hermes/__init__.py b/examples/hermes/__init__.py new file mode 100644 index 0000000000..acba72f988 --- /dev/null +++ b/examples/hermes/__init__.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Hermes example: an in-process per-session :class:`AgentRunnable`. + +Hosts the Hermes agent runtime that the Worker loads via ``agent_cls_path`` +(``examples.hermes.hermes.HermesAgent``). Unlike the OpenClaw +example (one Gateway subprocess per session), Hermes is a Python library, so +this runtime instantiates one Hermes ``AIAgent`` per RL session directly inside +the Worker process. +""" diff --git a/examples/hermes/config.yaml b/examples/hermes/config.yaml new file mode 100644 index 0000000000..91ca290c52 --- /dev/null +++ b/examples/hermes/config.yaml @@ -0,0 +1,164 @@ +experiment_name: online-rl +trial_name: trial0 + +seed: 1 +enable_offload: false +total_train_epochs: 10 +total_train_steps: 100 +tokenizer_path: ${actor.path} + +cluster: + n_nodes: 1 + n_gpus_per_node: 2 + fileroot: /tmp/areal/experiments + name_resolve: + type: nfs + nfs_record_root: /tmp/areal/name_resolve + +scheduler: + type: local + +rollout: + backend: "sglang:d1" + _version: v2 + experiment_name: ${experiment_name} + trial_name: ${trial_name} + max_concurrent_rollouts: 256 + queue_size: null + consumer_batch_size: ${train_dataset.batch_size} + max_head_offpolicyness: 2 + enable_rollout_tracing: false + scheduling_spec: ${actor.scheduling_spec} + fileroot: ${cluster.fileroot} + tokenizer_path: ${tokenizer_path} + dump_to_file: false + agent: + mode: online + tool_call_parser: qwen + reasoning_parser: qwen3 + export_style: individual + turn_discount: 1.0 + admin_api_key: sk-xxx + +gconfig: + n_samples: 1 + min_new_tokens: 0 + max_new_tokens: 2048 + greedy: false + temperature: 1.0 + +actor: + backend: "megatron:d1" + _version: v2 + admin_api_key: gsm8k-123 + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: Qwen/Qwen2.5-1.5B-Instruct + init_from_scratch: false + disable_dropout: true + gradient_checkpointing: true + dtype: bfloat16 + mb_spec: + max_tokens_per_mb: 15360 + optimizer: + type: adam + lr: 1.70e-5 + weight_decay: 0.017 + beta1: 0.9 + beta2: 0.999 + eps: 1e-8 + lr_scheduler_type: constant + gradient_clipping: 1.0 + warmup_steps_proportion: 0.001 + eps_clip: 0.4 + temperature: ${gconfig.temperature} + reward_scaling: 10.0 + reward_bias: -0.5 + kl_ctl: 0.0 + ppo_n_minibatches: 1 + recompute_logprob: true + use_decoupled_loss: true + rejection_sampling: + metric: ratio + upper: 5.0 + reward_norm: + mean_level: group + std_level: group + group_size: ${gconfig.n_samples} + adv_norm: + mean_level: batch + std_level: batch + weight_update_mode: xccl + max_new_tokens: ${gconfig.max_new_tokens} + scheduling_spec: + - task_type: worker + port_count: 2 + gpu: 1 + mem: 32 + cmd: python3 -m areal.infra.rpc.rpc_server + env_vars: {} + +# SGLang +sglang: + model_path: ${actor.path} + random_seed: ${seed} + skip_tokenizer_init: true + dtype: ${actor.dtype} + max_running_requests: null + context_length: 32768 + mem_fraction_static: 0.8 + +vllm: + model: ${actor.path} + seed: ${seed} + skip_tokenizer_init: false + dtype: ${actor.dtype} + max_model_len: 32768 + gpu_memory_utilization: 0.8 + +# datasets +train_dataset: + batch_size: 1 + path: "null" + type: rl + +# Utilities +saver: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +recover: + mode: disabled + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: 3600 + +evaluator: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +stats_logger: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + wandb: + mode: disabled + +perf_tracer: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + enabled: false + session_tracer: + enabled: false diff --git a/examples/hermes/hermes.py b/examples/hermes/hermes.py new file mode 100644 index 0000000000..f7ac4e8280 --- /dev/null +++ b/examples/hermes/hermes.py @@ -0,0 +1,444 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Hermes Agent for AReaL Agent Service (in-process per-session ``AIAgent``). + +Implements :class:`AgentRunnable` by instantiating **one Hermes ``AIAgent`` +per RL session** directly inside the Worker process. Unlike OpenClaw — which +exposes an OpenAI-compatible Gateway subprocess — Hermes +(``github.com/nousresearch/hermes-agent``) is a Python library whose +``AIAgent`` drives an OpenAI-compatible upstream itself. Each per-session +agent is therefore bound to its own upstream LLM (base URL + API key + model) +so that, during training, a session's turns can be attributed to a distinct +per-session key (``sk-sess-*``). + +Per turn the agent calls ``AIAgent.run_conversation`` with the conversation +history the DataProxy replays. Hermes rebuilds its message list from that +history on every call (it does not retain messages across calls), so the +replayed history stays the single source of truth and one cached agent can +safely serve every turn of a session. + +Requires the ``hermes-agent`` package importable as the top-level ``run_agent`` +module (``pip install hermes-agent``). + +Upstream selection +------------------ +Per turn the agent prefers the inference upstream the DataProxy injects via +``AgentRequest.metadata['areal_inference']`` (``base_url`` / ``api_key`` / +``model``), so a session's LLM calls flow through AReaL's inference service +under a per-session ``sk-sess-*`` key and get captured for training. Outside +training (e.g. the interactive demo) it falls back to the ``HERMES_UPSTREAM_*`` +environment variables. + +Environment variables +--------------------- + HERMES_UPSTREAM_BASE_URL — upstream LLM base URL (fallback when no + ``areal_inference`` metadata). + HERMES_UPSTREAM_API_KEY — upstream API key. + HERMES_UPSTREAM_MODEL — upstream model id (default ``default``). + HERMES_MAX_TURNS — max tool-calling iterations per turn (default 10). + HERMES_ENABLED_TOOLSETS — comma-separated toolsets to enable (optional). + HERMES_DISABLED_TOOLSETS — comma-separated toolsets to disable (optional). +""" + +from __future__ import annotations + +import asyncio +import json +import os +import time +import uuid +from dataclasses import dataclass, field +from typing import Any + +from areal.utils import logging +from areal.v2.agent_service.types import ( + AgentRequest, + AgentResponse, + EventEmitter, + StreamResponse, +) + +logger = logging.getLogger("HermesAgent") + + +def _split_csv(value: str) -> list[str] | None: + items = [item.strip() for item in value.split(",") if item.strip()] + return items or None + + +@dataclass +class _Upstream: + """Upstream LLM that a Hermes ``AIAgent`` routes its calls to.""" + + base_url: str + api_key: str + model: str + + @classmethod + def from_inference(cls, meta: dict[str, Any]) -> _Upstream | None: + """Build an upstream from a turn's ``metadata['areal_inference']``. + + The DataProxy injects this when a turn opts into AReaL's inference + service, so the session's LLM calls flow through the gateway under a + per-session ``sk-sess-*`` key and get captured for training. + """ + base_url = meta.get("base_url") or "" + api_key = meta.get("api_key") or "" + if not base_url or not api_key: + return None + model = meta.get("model") or "default" + return cls(base_url=base_url.rstrip("/"), api_key=api_key, model=model) + + @classmethod + def from_env(cls) -> _Upstream | None: + base_url = os.environ.get("HERMES_UPSTREAM_BASE_URL", "") + api_key = os.environ.get("HERMES_UPSTREAM_API_KEY", "") + if not base_url or not api_key: + return None + model = os.environ.get("HERMES_UPSTREAM_MODEL", "default") + return cls(base_url=base_url.rstrip("/"), api_key=api_key, model=model) + + +@dataclass +class _SessionState: + session_key: str + agent: Any # Hermes ``AIAgent`` instance. + upstream: _Upstream | None = None + metadata: dict[str, Any] = field(default_factory=dict) + + +class HermesAgent: + """AgentRunnable that runs one in-process Hermes ``AIAgent`` per session.""" + + def __init__(self, **_: Any) -> None: + self._max_turns = int(os.environ.get("HERMES_MAX_TURNS", "10")) + self._enabled_toolsets = _split_csv( + os.environ.get("HERMES_ENABLED_TOOLSETS", "") + ) + self._disabled_toolsets = _split_csv( + os.environ.get("HERMES_DISABLED_TOOLSETS", "") + ) + self._env_upstream = _Upstream.from_env() + + self._sessions: dict[str, _SessionState] = {} + self._locks: dict[str, asyncio.Lock] = {} + self._locks_guard = asyncio.Lock() + + logger.info( + "HermesAgent initialized (max_turns=%d, env_upstream=%s)", + self._max_turns, + self._env_upstream is not None, + ) + + # ------------------------------------------------------------------ + # AIAgent lifecycle + # ------------------------------------------------------------------ + + def _build_agent(self, upstream: _Upstream) -> Any: + """Construct a headless Hermes ``AIAgent`` bound to ``upstream``. + + Imported lazily so this module is importable (and unit-testable) + without the ``hermes-agent`` package installed. The persistence and + context-injection flags are disabled so the agent stays stateless + across turns — AReaL's DataProxy owns the conversation history. + """ + from run_agent import AIAgent # type: ignore[import-not-found] + + return AIAgent( + base_url=upstream.base_url, + api_key=upstream.api_key, + model=upstream.model, + max_iterations=self._max_turns, + enabled_toolsets=self._enabled_toolsets, + disabled_toolsets=self._disabled_toolsets, + save_trajectories=False, + quiet_mode=True, + skip_context_files=True, + skip_memory=True, + session_db=None, + ) + + async def _session_lock(self, session_key: str) -> asyncio.Lock: + async with self._locks_guard: + lock = self._locks.get(session_key) + if lock is None: + lock = asyncio.Lock() + self._locks[session_key] = lock + return lock + + async def _teardown_state(self, state: _SessionState) -> None: + # ``AIAgent`` cleanup is best-effort: not every build exposes a + # ``close``, and a failing teardown must not mask the caller's intent. + close = getattr(state.agent, "close", None) + if callable(close): + try: + result = close() + if asyncio.iscoroutine(result): + await result + except Exception: # noqa: BLE001 - teardown must not raise + logger.debug("AIAgent.close() failed", exc_info=True) + + async def _ensure_session( + self, + session_key: str, + upstream: _Upstream, + *, + rebuild: bool = False, + ) -> _SessionState: + lock = await self._session_lock(session_key) + async with lock: + existing = self._sessions.get(session_key) + if existing is not None: + if not rebuild: + return existing + await self._teardown_state(self._sessions.pop(session_key)) + # Building an AIAgent loads toolsets and is comparatively heavy; + # run it off the event loop so it does not block other sessions. + agent = await asyncio.to_thread(self._build_agent, upstream) + state = _SessionState( + session_key=session_key, agent=agent, upstream=upstream + ) + self._sessions[session_key] = state + logger.info( + "Built Hermes AIAgent (session=%s, model=%s)", + session_key, + upstream.model, + ) + return state + + # ------------------------------------------------------------------ + # Lifecycle hooks (optional members of the AgentRunnable protocol) + # ------------------------------------------------------------------ + + async def close_session(self, session_key: str) -> None: + state = self._sessions.pop(session_key, None) + if state is not None: + await self._teardown_state(state) + # Drop the per-session lock so ``_locks`` does not grow unbounded as + # sessions are created and destroyed over a long-running worker. + async with self._locks_guard: + self._locks.pop(session_key, None) + + async def close_all_sessions(self) -> None: + for key in list(self._sessions.keys()): + await self.close_session(key) + + # ------------------------------------------------------------------ + # Per-turn execution + # ------------------------------------------------------------------ + + async def _resolve_state(self, request: AgentRequest) -> _SessionState: + """Pick (or build) the Hermes ``AIAgent`` backing this turn. + + Session-first: an existing per-session agent is reused as-is, so a turn + need not re-supply its upstream once the session is live. When no agent + exists yet, an upstream must be resolvable — preferring the per-session + inference routing the DataProxy injected (self-evolution, calls flow + through AReaL's inference service under a ``sk-sess-*`` key), else the + process-wide env upstream (e.g. the interactive demo). If a live + session's upstream changes (env → AReaL inference), rebuild so the agent + routes to the new endpoint/key. + """ + inf_meta = (request.metadata or {}).get("areal_inference") + upstream = ( + _Upstream.from_inference(inf_meta) if inf_meta else None + ) or self._env_upstream + + state = self._sessions.get(request.session_key) + if state is None: + if upstream is None: + raise RuntimeError( + f"No agent for session '{request.session_key}': send the turn " + "with inference-routing fields ('inf_base_url' + " + "'session_api_key', so the DataProxy injects an inference " + "upstream) or set the HERMES_UPSTREAM_* env." + ) + return await self._ensure_session(request.session_key, upstream) + if upstream is not None and state.upstream != upstream: + return await self._ensure_session( + request.session_key, upstream, rebuild=True + ) + return state + + async def _run_conversation( + self, + state: _SessionState, + message: str, + history: list[dict[str, Any]], + ) -> dict[str, Any]: + """Run one blocking Hermes turn off the event loop, serialized per session. + + ``run_conversation`` is synchronous and mutates agent state, so the + per-session lock keeps a session's turns from overlapping. The full + replayed history is passed every turn (Hermes rebuilds its message list + from it), keeping the DataProxy's history the single source of truth. + """ + lock = await self._session_lock(state.session_key) + async with lock: + return await asyncio.to_thread( + state.agent.run_conversation, + message, + conversation_history=history or None, + task_id=f"areal-{uuid.uuid4().hex}", + ) + + @staticmethod + def _collect_tool_calls(result: dict[str, Any]) -> list[dict[str, Any]]: + """Surface tool calls from the result transcript for observability. + + Like the OpenClaw runtime, tool calls are reported in ``metadata`` only + and deliberately NOT emitted as ``tool_call`` events: Hermes executes + tools internally, so emitting a ``tool_call`` without the framework also + feeding a paired ``tool_result`` would make the DataProxy build an + invalid chat-completions history that the upstream rejects on replay. + """ + tool_calls: list[dict[str, Any]] = [] + for msg in result.get("messages") or []: + if msg.get("role") != "assistant": + continue + for call in msg.get("tool_calls") or []: + fn = call.get("function") or {} + tool_calls.append( + { + "name": fn.get("name", ""), + "input": fn.get("arguments", ""), + } + ) + return tool_calls + + async def run( + self, + request: AgentRequest, + *, + emitter: EventEmitter, + ) -> AgentResponse | StreamResponse: + """Single entry point for both the structured and chat-completions channels. + + The ``/v1/chat/completions`` bridge stashes the original request body + under ``metadata['chat_request']``; its presence selects the + chat-completions channel (return a :class:`StreamResponse` carrying an + OpenAI-compatible response synthesized from Hermes' final text). + Otherwise the structured channel runs: report the final text through + ``emitter`` and return an :class:`AgentResponse`. + """ + state = await self._resolve_state(request) + + chat_request = (request.metadata or {}).get("chat_request") + if chat_request is not None: + return await self._chat_completions(state, dict(chat_request)) + + result = await self._run_conversation(state, request.message, request.history) + final = result.get("final_response") or "" + if final: + await emitter.emit_delta(final) + + return AgentResponse( + summary=final, + metadata={ + "tool_calls": self._collect_tool_calls(result), + "completed": bool(result.get("completed")), + "api_calls": result.get("api_calls", 0), + }, + ) + + # ------------------------------------------------------------------ + # Chat-completions channel + # ------------------------------------------------------------------ + + @staticmethod + def _extract_chat(chat_body: dict[str, Any]) -> tuple[str, list[dict[str, Any]]]: + """Split a chat-completions body into (last user message, prior history).""" + messages = list(chat_body.get("messages") or []) + message = "" + for idx in range(len(messages) - 1, -1, -1): + if messages[idx].get("role") == "user": + content = messages[idx].get("content", "") + if isinstance(content, list): + content = "".join( + block.get("text", "") + for block in content + if isinstance(block, dict) and block.get("type") == "text" + ) + message = content or "" + history = messages[:idx] + return message, history + return message, messages + + async def _chat_completions( + self, + state: _SessionState, + chat_body: dict[str, Any], + ) -> StreamResponse: + """Chat-completions channel for the ``/v1/chat/completions`` bridge. + + Hermes has no OpenAI-compatible HTTP server to relay byte-for-byte, so + this runs the structured turn and re-encodes Hermes' final text into the + OpenAI chat-completions wire format (SSE when ``stream`` is set, a single + JSON object otherwise), returned as a :class:`StreamResponse` the Worker + and DataProxy relay verbatim. + """ + message, history = self._extract_chat(chat_body) + result = await self._run_conversation(state, message, history) + final = result.get("final_response") or "" + model = chat_body.get("model") or ( + state.upstream.model if state.upstream else "hermes-agent" + ) + stream = bool(chat_body.get("stream")) + + if stream: + return self._chat_sse_response(final, model) + return self._chat_json_response(final, model) + + @staticmethod + def _chat_json_response(content: str, model: str) -> StreamResponse: + payload = { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": int(time.time()), + "model": model, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": content}, + "finish_reason": "stop", + } + ], + } + body = json.dumps(payload).encode("utf-8") + + async def _iter() -> Any: + yield body + + return StreamResponse( + status_code=200, + headers={"content-type": "application/json"}, + body=_iter(), + ) + + @staticmethod + def _chat_sse_response(content: str, model: str) -> StreamResponse: + cmpl_id = f"chatcmpl-{uuid.uuid4().hex}" + created = int(time.time()) + + def _chunk(delta: dict[str, Any], finish: str | None) -> bytes: + payload = { + "id": cmpl_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish}], + } + return f"data: {json.dumps(payload)}\n\n".encode() + + async def _iter() -> Any: + yield _chunk({"role": "assistant"}, None) + if content: + yield _chunk({"content": content}, None) + yield _chunk({}, "stop") + yield b"data: [DONE]\n\n" + + return StreamResponse( + status_code=200, + headers={"content-type": "text/event-stream"}, + body=_iter(), + ) diff --git a/examples/hermes/hermes_loop.py b/examples/hermes/hermes_loop.py new file mode 100644 index 0000000000..6c764efd7c --- /dev/null +++ b/examples/hermes/hermes_loop.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Interactive conversation loop against a running Hermes Agent Service. + +This is the standalone chat window: it assumes the Agent Service is **already +running** (started separately with ``areal agent run``) and a per-session +``sk-sess-*`` key has **already been minted** on the inference gateway (with +``start_session.py``). It only opens a ``You:`` prompt and drives the agent. + +Each user message becomes one turn of the Hermes conversation, posted to the +agent gateway's ``/v1/responses``. When ``--inf-base-url`` and +``--session-api-key`` are supplied, every turn forwards the inference-routing +fields so the agent's LLM calls flow through AReaL's inference service under +that key and the trajectory is captured for training; score it afterwards with +``set_reward.py``. + +Usage:: + + # plain chat (env upstream; agent admin key authorizes the gateway) + python hermes_loop.py http:// --admin-api-key + + # self-evolution: capture the trajectory under a session key + python hermes_loop.py http:// --admin-api-key \\ + --inf-base-url http:// --session-api-key sk-sess-xxxx +""" + +from __future__ import annotations + +import argparse +import asyncio +import time + +import httpx + + +async def interactive_loop( + gateway_addr: str, + admin_key: str, + inference: dict[str, str] | None = None, +) -> None: + session_key = f"hermes-{int(time.time())}" + print(f"Session: {session_key}") + if inference: + print(f"Self-evolution: routing LLM calls via {inference['inf_base_url']}") + print("Type your message (or 'quit' to exit):\n") + + async with httpx.AsyncClient(timeout=120.0) as client: + while True: + try: + user_input = input("You: ") + except (EOFError, KeyboardInterrupt): + break + if user_input.strip().lower() in ("quit", "exit", "q"): + break + if not user_input.strip(): + continue + + payload: dict[str, object] = { + "input": [{"type": "message", "content": user_input}], + "model": "hermes-agent", + "user": session_key, + } + if inference: + payload.update(inference) + + resp = await client.post( + f"{gateway_addr}/v1/responses", + json=payload, + headers={"Authorization": f"Bearer {admin_key}"}, + ) + data = resp.json() + + if data.get("status") == "completed": + for item in data.get("output", []): + if item.get("type") == "message": + for block in item.get("content", []): + if block.get("type") == "output_text": + print(f"Agent: {block['text']}") + elif item.get("type") == "function_call": + print(f"[tool] {item.get('name', '')}") + print() + elif data.get("error"): + print(f"Error: {data['error'].get('message', '')[:200]}\n") + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Interactive chat against a running Hermes Agent Service" + ) + parser.add_argument( + "gateway_addr", + help="Agent Service gateway URL (from 'areal agent status')", + ) + parser.add_argument( + "--admin-api-key", + required=True, + help="Agent Service admin key (the --admin-api-key passed to 'areal agent run')", + ) + parser.add_argument( + "--inf-base-url", + default="", + help="Inference gateway base URL (enables self-evolution capture)", + ) + parser.add_argument( + "--inf-model", + default="", + help="Model id the agent requests from the inference service", + ) + parser.add_argument( + "--session-api-key", + default="", + help="Per-session sk-sess-* key from start_session.py (required for capture)", + ) + args = parser.parse_args() + + gateway_addr = args.gateway_addr.rstrip("/") + + inference: dict[str, str] | None = None + if args.inf_base_url or args.session_api_key: + if not (args.inf_base_url and args.session_api_key): + raise SystemExit( + "self-evolution requires BOTH --inf-base-url and --session-api-key" + ) + inference = { + "inf_base_url": args.inf_base_url.rstrip("/"), + "inf_model": args.inf_model, + "session_api_key": args.session_api_key, + } + + asyncio.run( + interactive_loop( + gateway_addr, admin_key=args.admin_api_key, inference=inference + ) + ) + + +if __name__ == "__main__": + main() diff --git a/examples/hermes/set_reward.py b/examples/hermes/set_reward.py new file mode 100644 index 0000000000..6c5b3a70b9 --- /dev/null +++ b/examples/hermes/set_reward.py @@ -0,0 +1,140 @@ +#!/usr/bin/env python3 +"""Set reward for the last interaction in an active session. + +Requires the per-session ``sk-sess-*`` key minted by ``start_session.py``. + +Usage: + python set_reward.py http://host:port --api-key --reward 1.0 + python set_reward.py http://host:port --api-key --reward 0.0 --interaction-id cmpl_abc +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys + +import requests + +# ============================================================================= +# CLI formatting helpers (ANSI colors + output helpers) +# ============================================================================= +BOLD = "\033[1m" +DIM = "\033[2m" +RESET = "\033[0m" +GREEN = "\033[0;32m" +CYAN = "\033[0;36m" +RED = "\033[0;31m" +BLUE = "\033[0;34m" + + +def info(msg: str) -> None: + print(f" {CYAN}ℹ{RESET} {msg}") + + +def success(msg: str) -> None: + print(f" {GREEN}✔{RESET} {msg}") + + +def error(msg: str) -> None: + print(f" {RED}✘{RESET} {msg}") + + +def dim(msg: str) -> None: + print(f" {DIM}{msg}{RESET}") + + +def die(msg: str) -> None: + error(msg) + sys.exit(1) + + +def show_request(method: str, path: str, auth_label: str, gateway_url: str) -> None: + print(f" {DIM}{method} {gateway_url}/{path}{RESET}") + print(f" {DIM}Auth: {auth_label}{RESET}") + + +def show_response(status_code: int, body: str) -> None: + if 200 <= status_code < 300: + print(f" {GREEN}HTTP {status_code}{RESET}") + else: + print(f" {RED}HTTP {status_code}{RESET}") + if body: + try: + formatted = json.dumps(json.loads(body), indent=2) + for line in formatted.split("\n"): + print(f" {DIM}{line}{RESET}") + except (json.JSONDecodeError, ValueError): + for line in body.split("\n"): + print(f" {DIM}{line}{RESET}") + + +def header(title: str) -> None: + """Print a boxed header.""" + print() + print( + f"{BOLD}{BLUE}══════════════════════════════════════════════════════════════{RESET}" + ) + print(f"{BOLD}{BLUE} {title}{RESET}") + print( + f"{BOLD}{BLUE}══════════════════════════════════════════════════════════════{RESET}" + ) + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Set reward for an interaction in an AReaL RL session" + ) + parser.add_argument("gateway_url", help="Proxy gateway URL") + parser.add_argument( + "--api-key", + required=True, + help="Per-session sk-sess-* key from start_session.py", + ) + parser.add_argument( + "--reward", + type=float, + default=float(os.getenv("REWARD", "1.0")), + help="Reward value (default: 1.0, env: REWARD)", + ) + parser.add_argument( + "--interaction-id", + default=None, + help="Specific interaction ID (default: last interaction)", + ) + args = parser.parse_args() + + header("Set Reward") + info(f"Assigning reward={args.reward} to the last interaction.") + show_request("POST", "rl/set_reward", "Bearer ***", args.gateway_url) + + reward_body: dict = {"reward": args.reward} + if args.interaction_id is not None: + reward_body["interaction_id"] = args.interaction_id + + dim(f"Request body: {json.dumps(reward_body)}") + + try: + resp = requests.post( + f"{args.gateway_url}/rl/set_reward", + headers={ + "Content-Type": "application/json", + "Authorization": f"Bearer {args.api_key}", + }, + json=reward_body, + timeout=10, + ) + except requests.RequestException as e: + die(f"Failed to reach gateway: {e}") + + show_response(resp.status_code, resp.text) + + if resp.status_code != 200: + die(f"set_reward failed (HTTP {resp.status_code}).") + + success(f"Reward {args.reward} applied") + + +if __name__ == "__main__": + main() diff --git a/examples/hermes/start_session.py b/examples/hermes/start_session.py new file mode 100644 index 0000000000..ff29a7fb66 --- /dev/null +++ b/examples/hermes/start_session.py @@ -0,0 +1,198 @@ +#!/usr/bin/env python3 +"""Start an RL session on AReaL's inference gateway. + +Mints a per-session ``sk-sess-*`` key on the inference gateway that +``train.py`` embeds, then prints it. You forward that key to the Hermes +Agent Service (``hermes_loop.py``) so the agent's LLM calls flow through the +inference service under it and get captured as a training trajectory; you then +score it with ``set_reward.py``. + +The inference gateway's admin key is whatever you passed as +``rollout.admin_api_key`` to ``train.py``. + +Pass ``--api-key`` with a previously issued key to **refresh** an existing +session: the old session is ended (default reward 0 if none was set), its +trajectory is exported, and a new session is started reusing the same key. + +Usage: + python start_session.py http://host:port --admin-key sk-xxx + python start_session.py http://host:port --admin-key sk-xxx --task-id my-task + python start_session.py http://host:port --admin-key sk-xxx --api-key +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys + +import requests + +# ============================================================================= +# CLI formatting helpers (ANSI colors + output helpers) +# ============================================================================= +BOLD = "\033[1m" +DIM = "\033[2m" +RESET = "\033[0m" +GREEN = "\033[0;32m" +CYAN = "\033[0;36m" +YELLOW = "\033[0;33m" +RED = "\033[0;31m" +BLUE = "\033[0;34m" + + +def info(msg: str) -> None: + print(f" {CYAN}ℹ{RESET} {msg}") + + +def success(msg: str) -> None: + print(f" {GREEN}✔{RESET} {msg}") + + +def error(msg: str) -> None: + print(f" {RED}✘{RESET} {msg}") + + +def arrow(msg: str) -> None: + print(f" {YELLOW}→{RESET} {msg}") + + +def die(msg: str) -> None: + error(msg) + sys.exit(1) + + +def show_request(method: str, path: str, auth_label: str, gateway_url: str) -> None: + print(f" {DIM}{method} {gateway_url}/{path}{RESET}") + print(f" {DIM}Auth: {auth_label}{RESET}") + + +def show_response(status_code: int, body: str) -> None: + if 200 <= status_code < 300: + print(f" {GREEN}HTTP {status_code}{RESET}") + else: + print(f" {RED}HTTP {status_code}{RESET}") + if body: + try: + formatted = json.dumps(json.loads(body), indent=2) + for line in formatted.split("\n"): + print(f" {DIM}{line}{RESET}") + except (json.JSONDecodeError, ValueError): + for line in body.split("\n"): + print(f" {DIM}{line}{RESET}") + + +def header(title: str) -> None: + """Print a boxed header.""" + print() + print( + f"{BOLD}{BLUE}══════════════════════════════════════════════════════════════{RESET}" + ) + print(f"{BOLD}{BLUE} {title}{RESET}") + print( + f"{BOLD}{BLUE}══════════════════════════════════════════════════════════════{RESET}" + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Start an AReaL RL session") + parser.add_argument("gateway_url", help="Inference gateway URL") + parser.add_argument( + "--admin-key", + default=os.getenv("ADMIN_KEY", "areal-admin-key"), + help="Inference gateway admin key (== rollout.admin_api_key; env: ADMIN_KEY)", + ) + parser.add_argument( + "--task-id", + default=os.getenv("TASK_ID", "demo-task"), + help="Task identifier (env: TASK_ID)", + ) + parser.add_argument( + "--api-key", + default=os.getenv("SESSION_API_KEY"), + help="Reuse a previously issued API key (refresh). (env: SESSION_API_KEY)", + ) + args = parser.parse_args() + + is_refresh = args.api_key is not None + header("Refresh Session" if is_refresh else "Start Session") + if is_refresh: + info( + "Refreshing: end old session → export trajectory → start new session (same key)" + ) + else: + info("Requesting a new RL session (admin auth → gateway routes to a worker)") + show_request("POST", "rl/start_session", "Bearer ***", args.gateway_url) + + try: + resp = requests.post( + f"{args.gateway_url}/rl/start_session", + headers={ + "Content-Type": "application/json", + "Authorization": f"Bearer {args.admin_key}", + }, + json={"task_id": args.task_id, "api_key": args.api_key}, + timeout=130 if is_refresh else 10, + ) + except requests.RequestException as e: + die(f"Failed to reach gateway: {e}") + + show_response(resp.status_code, resp.text) + + if resp.status_code != 201: + if resp.status_code == 429 and is_refresh: + die( + "Refresh timed out — the training pipeline hasn't cycled yet. " + "Retry in a few seconds." + ) + die( + "start_session failed. " + "If HTTP 429, no capacity — the RL trainer hasn't granted capacity yet." + ) + + # v2 inference gateway returns 201 with a flat list of session credentials: + # {"group_id": ..., "sessions": [{"session_id": ..., "session_api_key": ...}]} + try: + data = resp.json() + session = data["sessions"][0] + session_api_key = session["session_api_key"] + session_id = session["session_id"] + except (ValueError, KeyError, IndexError) as e: + die(f"Failed to parse response: {e}") + + success("Session started!") + arrow(f"Session ID : {BOLD}{session_id}{RESET}") + arrow(f"API Key : {BOLD}{session_api_key}{RESET}") + print() + info("Forward this key to the Hermes Agent Service to capture the trajectory:") + print() + print( + f" python hermes_loop.py " + f" --admin-api-key " + f" --inf-base-url {args.gateway_url}" + f" --session-api-key {session_api_key}" + ) + print() + info("Then score the episode with the same key:") + print() + print( + f" python set_reward.py {args.gateway_url}" + f" --api-key {session_api_key} --reward 1.0" + ) + print() + info("To start the next episode reusing the same key:") + print() + print( + f" python start_session.py {args.gateway_url}" + f" --admin-key {args.admin_key} --api-key {session_api_key}" + ) + print() + + # Machine-readable output on stderr for scripting + print(f"SESSION_API_KEY={session_api_key}", file=sys.stderr) + print(f"SESSION_ID={session_id}", file=sys.stderr) + + +if __name__ == "__main__": + main() diff --git a/examples/hermes/train.py b/examples/hermes/train.py new file mode 100644 index 0000000000..3c229838a5 --- /dev/null +++ b/examples/hermes/train.py @@ -0,0 +1,17 @@ +"""Online training: external users interact via proxy gateway URL.""" + +import sys + +from areal import PPOTrainer +from areal.api.cli_args import PPOConfig, load_expr_config + + +def main(args): + config, _ = load_expr_config(args, PPOConfig) + + with PPOTrainer(config) as trainer: + trainer.train() + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/examples/math/gsm8k_icepop.yaml b/examples/math/gsm8k_icepop.yaml new file mode 100644 index 0000000000..6209530d1c --- /dev/null +++ b/examples/math/gsm8k_icepop.yaml @@ -0,0 +1,190 @@ +experiment_name: gsm8k-icepop +trial_name: trial0 + +seed: 1 +enable_offload: false +total_train_epochs: 10 +tokenizer_path: ${actor.path} + +cluster: + n_nodes: 1 + n_gpus_per_node: 8 + fileroot: /tmp/areal/experiments + name_resolve: + type: nfs + nfs_record_root: /tmp/areal/name_resolve + + +scheduler: + type: null + +rollout: + backend: "sglang:d4p1t1" + experiment_name: ${experiment_name} + trial_name: ${trial_name} + max_concurrent_rollouts: 256 + queue_size: null + consumer_batch_size: ${train_dataset.batch_size} + max_head_offpolicyness: 2 + enable_rollout_tracing: false + scheduling_spec: ${actor.scheduling_spec} + fileroot: ${cluster.fileroot} + tokenizer_path: ${tokenizer_path} + dump_to_file: false + agent: + mode: inline + export_style: individual + turn_discount: 1.0 + +gconfig: + n_samples: 4 + min_new_tokens: 0 + max_new_tokens: 1024 + max_tokens: 2048 + greedy: false + temperature: 1.0 + +actor: + backend: "fsdp:d4p1t1" + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: Qwen/Qwen2.5-1.5B-Instruct + init_from_scratch: false + disable_dropout: true + gradient_checkpointing: true + dtype: bfloat16 + mb_spec: + max_tokens_per_mb: 10240 + packing_algorithm: ffd + optimizer: + type: adam + lr: 1.70e-5 + weight_decay: 0.017 + beta1: 0.9 + beta2: 0.999 + eps: 1e-8 + lr_scheduler_type: constant + gradient_clipping: 1.0 + warmup_steps_proportion: 0.001 + eps_clip: 0.4 + temperature: ${gconfig.temperature} + reward_scaling: 10.0 + reward_bias: -0.5 + kl_ctl: 0.0 + ppo_n_minibatches: 1 + use_decoupled_loss: true + rejection_sampling: + level: token + action: mask + metric: ratio + lower: 0.5 + upper: 5.0 + reward_norm: + mean_level: group + std_level: group + group_size: ${gconfig.n_samples} + adv_norm: + mean_level: batch + std_level: batch + weight_update_mode: xccl + max_new_tokens: ${gconfig.max_new_tokens} + scheduling_spec: + - task_type: worker + port_count: 2 + gpu: 1 + mem: 32 + cmd: python3 -m areal.infra.rpc.rpc_server + env_vars: {} + +ref: + backend: ${actor.backend} + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: ${actor.path} + init_from_scratch: false + disable_dropout: true + dtype: ${actor.dtype} + mb_spec: + max_tokens_per_mb: 10240 + packing_algorithm: ffd + optimizer: null + scheduling_strategy: + type: colocation + target: actor + scheduling_spec: ${actor.scheduling_spec} + +# SGLang +sglang: + model_path: ${actor.path} + random_seed: ${seed} + skip_tokenizer_init: true + dtype: ${actor.dtype} + max_running_requests: null + context_length: 32768 + mem_fraction_static: 0.8 + +vllm: + model: ${actor.path} + seed: ${seed} + skip_tokenizer_init: false + dtype: ${actor.dtype} + max_model_len: 32768 + gpu_memory_utilization: 0.8 + +# datasets +train_dataset: + batch_size: 256 + shuffle: true + pin_memory: true + num_workers: 4 + path: openai/gsm8k + type: rl + max_length: 1024 + +valid_dataset: + batch_size: 256 + pin_memory: true + num_workers: 4 + path: openai/gsm8k + type: rl + +# Utilities +saver: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +recover: + mode: disabled + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: 3600 + +evaluator: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +stats_logger: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + wandb: + mode: disabled + +perf_tracer: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + enabled: false + session_tracer: + enabled: false diff --git a/examples/math/gsm8k_kpop.yaml b/examples/math/gsm8k_kpop.yaml new file mode 100644 index 0000000000..2a0397fe26 --- /dev/null +++ b/examples/math/gsm8k_kpop.yaml @@ -0,0 +1,189 @@ +experiment_name: gsm8k-kpop +trial_name: trial0 + +seed: 1 +enable_offload: false +total_train_epochs: 10 +tokenizer_path: ${actor.path} + +cluster: + n_nodes: 1 + n_gpus_per_node: 8 + fileroot: /tmp/areal/experiments + name_resolve: + type: nfs + nfs_record_root: /tmp/areal/name_resolve + + +scheduler: + type: null + +rollout: + backend: "sglang:d4p1t1" + experiment_name: ${experiment_name} + trial_name: ${trial_name} + max_concurrent_rollouts: 256 + queue_size: null + consumer_batch_size: ${train_dataset.batch_size} + max_head_offpolicyness: 2 + enable_rollout_tracing: false + scheduling_spec: ${actor.scheduling_spec} + fileroot: ${cluster.fileroot} + tokenizer_path: ${tokenizer_path} + dump_to_file: false + agent: + mode: inline + export_style: individual + turn_discount: 1.0 + +gconfig: + n_samples: 4 + min_new_tokens: 0 + max_new_tokens: 1024 + max_tokens: 2048 + greedy: false + temperature: 1.0 + +actor: + backend: "fsdp:d4p1t1" + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: Qwen/Qwen2.5-1.5B-Instruct + init_from_scratch: false + disable_dropout: true + gradient_checkpointing: true + dtype: bfloat16 + mb_spec: + max_tokens_per_mb: 10240 + packing_algorithm: ffd + optimizer: + type: adam + lr: 1.70e-5 + weight_decay: 0.017 + beta1: 0.9 + beta2: 0.999 + eps: 1e-8 + lr_scheduler_type: constant + gradient_clipping: 1.0 + warmup_steps_proportion: 0.001 + eps_clip: 0.4 + temperature: ${gconfig.temperature} + reward_scaling: 10.0 + reward_bias: -0.5 + kl_ctl: 0.0 + ppo_n_minibatches: 1 + use_decoupled_loss: true + rejection_sampling: + level: token + action: mask + metric: binary_kl + upper: 2.0 + reward_norm: + mean_level: group + std_level: group + group_size: ${gconfig.n_samples} + adv_norm: + mean_level: batch + std_level: batch + weight_update_mode: xccl + max_new_tokens: ${gconfig.max_new_tokens} + scheduling_spec: + - task_type: worker + port_count: 2 + gpu: 1 + mem: 32 + cmd: python3 -m areal.infra.rpc.rpc_server + env_vars: {} + +ref: + backend: ${actor.backend} + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: ${actor.path} + init_from_scratch: false + disable_dropout: true + dtype: ${actor.dtype} + mb_spec: + max_tokens_per_mb: 10240 + packing_algorithm: ffd + optimizer: null + scheduling_strategy: + type: colocation + target: actor + scheduling_spec: ${actor.scheduling_spec} + +# SGLang +sglang: + model_path: ${actor.path} + random_seed: ${seed} + skip_tokenizer_init: true + dtype: ${actor.dtype} + max_running_requests: null + context_length: 32768 + mem_fraction_static: 0.8 + +vllm: + model: ${actor.path} + seed: ${seed} + skip_tokenizer_init: false + dtype: ${actor.dtype} + max_model_len: 32768 + gpu_memory_utilization: 0.8 + +# datasets +train_dataset: + batch_size: 256 + shuffle: true + pin_memory: true + num_workers: 4 + path: openai/gsm8k + type: rl + max_length: 1024 + +valid_dataset: + batch_size: 256 + pin_memory: true + num_workers: 4 + path: openai/gsm8k + type: rl + +# Utilities +saver: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +recover: + mode: disabled + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: 3600 + +evaluator: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +stats_logger: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + wandb: + mode: disabled + +perf_tracer: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + enabled: false + session_tracer: + enabled: false diff --git a/examples/skypilot/README.md b/examples/skypilot/README.md index c358d0634e..74e152acfa 100644 --- a/examples/skypilot/README.md +++ b/examples/skypilot/README.md @@ -25,7 +25,7 @@ resources: cpus: 8+ memory: 32GB+ disk_size: 256GB - image_id: docker:ghcr.io/areal-project/areal-runtime:v1.0.4-sglang + image_id: docker:ghcr.io/areal-project/areal-runtime:v2.0.0-sglang num_nodes: 1 @@ -78,7 +78,7 @@ Specify the resources and image used to run the experiment. ```yaml resources: accelerators: A100:8 - image_id: docker:ghcr.io/areal-project/areal-runtime:v1.0.4-sglang + image_id: docker:ghcr.io/areal-project/areal-runtime:v2.0.0-sglang memory: 256+ cpus: 32+ diff --git a/examples/skypilot/ray_cluster.sky.yaml b/examples/skypilot/ray_cluster.sky.yaml index 341f7c7eeb..e79f838397 100644 --- a/examples/skypilot/ray_cluster.sky.yaml +++ b/examples/skypilot/ray_cluster.sky.yaml @@ -1,7 +1,7 @@ resources: accelerators: A100:8 - image_id: docker:ghcr.io/areal-project/areal-runtime:v1.0.4-sglang + image_id: docker:ghcr.io/areal-project/areal-runtime:v2.0.0-sglang memory: 32+ cpus: 8+ diff --git a/examples/skypilot/single_node.sky.yaml b/examples/skypilot/single_node.sky.yaml index 1bd90b8518..bf4c7f29cf 100644 --- a/examples/skypilot/single_node.sky.yaml +++ b/examples/skypilot/single_node.sky.yaml @@ -8,7 +8,7 @@ resources: cpus: 8+ memory: 32GB+ disk_size: 256GB - image_id: docker:ghcr.io/areal-project/areal-runtime:v1.0.4-sglang + image_id: docker:ghcr.io/areal-project/areal-runtime:v2.0.0-sglang num_nodes: 1 diff --git a/examples/swe/README.md b/examples/swe/README.md new file mode 100644 index 0000000000..ca3985ec40 --- /dev/null +++ b/examples/swe/README.md @@ -0,0 +1,190 @@ +# SWE-bench RL training with AReaL-SWEAgent + +This example runs SWE-bench coding-agent RL (GRPO) in AReaL. The actual agent loop, +sandboxing and reward computation live in a **separate repository**, +[AReaL-SWEAgent](https://github.com/areal-project/AReaL-SWEAgent): for each SWE-bench +instance it edits code inside an isolated sandbox, grades the patch, and returns a +reward for the RL loop. AReaL serves the policy being trained and drives rollouts +through its OpenAI-compatible proxy. + +``` +AReaL (this repo) AReaL-SWEAgent (separate checkout) + PPOTrainer / rollout proxy ──▶ run_agent_with_reward() ──▶ sandbox + grade + ▲ │ + └──────────────── reward, stats ◀──────────────────────────────┘ +``` + +## 1. Get AReaL-SWEAgent + +Clone it next to AReaL (the default lookup path is `../AReaL-SWEAgent`) and install its +dependencies into the same environment AReaL workers run in: + +```bash +git clone https://github.com/areal-project/AReaL-SWEAgent.git +cd AReaL-SWEAgent +pip install -r requirements.txt +``` + +AReaL imports it as the Python package `aweagent`; the repository directory name +(`AReaL-SWEAgent`) and the package name (`aweagent`) intentionally differ. + +## 2. Point AReaL at the checkout + +AReaL discovers the checkout in this order (first match wins): + +1. `econfig.agent_root` in your YAML (legacy alias: `econfig.swe_agent_root`). +1. `AWEAGENT_ROOT` / `SWE_AGENT_ROOT` environment variable. +1. `../AReaL-SWEAgent` relative to the AReaL repo root. + +In a multi-node / containerized setup the checkout must be importable on every worker, +so put it on shared storage and set both the env var and `PYTHONPATH`, e.g. in the +worker `scheduling_spec.env_vars`: + +```yaml +env_vars: + AWEAGENT_ROOT: "/path/to/AReaL-SWEAgent" + SWE_AGENT_ROOT: "/path/to/AReaL-SWEAgent" + PYTHONPATH: "/path/to/AReaL-SWEAgent:/path/to/AReaL" + # URL of the AEnvironment sandbox service used by AReaL-SWEAgent. + AENV_SYSTEM_URL: "http://your-aenv-service:8080" +``` + +## 3. Configure the agent (econfig) + +The SWE example reads an `econfig` block (see `examples/swe/utils.py`): + +| Field | Meaning | +| ----------------------- | ----------------------------------------------------------------------------------------------------------- | +| `agent_type` | Agent to run: `swe` (built-in tool-use) or `cc` (Claude Code). | +| `agent_config` | Generic config name under `AReaL-SWEAgent/aweagent/configs/`. Overrides the per-type fields below when set. | +| `swe_agent_config` | Config used when `agent_type=swe` (default `1_0_0/min-swe-agent-train-top1`). | +| `cc_agent_config` | Config used when `agent_type=cc`. | +| `agent_root` | Path to the AReaL-SWEAgent checkout (see section 2). | +| `step_limit` | Max agent interaction steps per episode. | +| `max_completion_tokens` | Max completion tokens per agent LLM call. | +| `timeout` | Max wall-clock time per episode (seconds). | + +## 4. Dataset format + +`train_swe_rl.py` loads a JSONL file (`train_dataset.path`) where each line is a +SWE-bench instance with at least: + +- `instance_id` — e.g. `django__django-10097` +- `problem_statement` — the GitHub issue text +- `eval_script` — shell script used by AReaL-SWEAgent to grade the patch + +## 5. Run + +Edit the placeholder paths in `qwen3_30b_a3b_grpo.yaml` (model, dataset, +`AWEAGENT_ROOT`, caches, sandbox/W&B endpoints), then launch: + +```bash +python -m examples.swe.train_swe_rl --config examples/swe/qwen3_30b_a3b_grpo.yaml +``` + +With `scheduler.type=slurm`, AReaL launches the rollout / actor / proxy workers; each +rollout calls into AReaL-SWEAgent, which runs the agent in a sandbox and returns the +reward. + +## 6. Set up the AEnvironment backend + +`AENV_SYSTEM_URL` (section 2) must point at a running +[AEnvironment](https://github.com/inclusionAI/AEnvironment) platform. AReaL never talks +to a sandbox directly — AReaL-SWEAgent asks the AEnvironment API service to spin up two +kinds of environment per rollout: + +| Role | Environment AReaL-SWEAgent requests | Ships in the AEnvironment repo? | +| ----------------------------------------------------------------------- | ------------------------------------------------------------ | --------------------------------------------- | +| **Task sandbox** — where the agent runs shell commands and edits code | `persistent-bash-env` (swe agent) / `cc-bash-env` (cc agent) | No — provided by your AEnvironment deployment | +| **Grading sandbox** — where the patch is scored resolved / not-resolved | `swebench` (overridable via `SWEBENCH_EVAL_ENV`) | Yes — `aenv/builtin-envs/swebench` | + +### 6.1 Deploy the platform + +AEnvironment deploys on Kubernetes via Helm (see its +[deployment guide](https://github.com/inclusionAI/AEnvironment/blob/main/docs/getting_started/deployment.md)): + +```bash +git clone https://github.com/inclusionAI/AEnvironment.git +cd AEnvironment +helm install aenv-platform ./deploy \ + --namespace aenv --create-namespace --wait --timeout 10m +``` + +The API service is then reachable in-cluster at port `8080` — this is the value for +`AENV_SYSTEM_URL`: + +``` +AENV_SYSTEM_URL = http://api-service.aenv.svc.cluster.local:8080 +``` + +### 6.2 Register the `swebench` grading environment (guaranteed example) + +The grading environment is the reward source and is the one piece shipped in the +AEnvironment repo, so it always works out of the box. Install the CLI, point it at your +deployment, then build and publish the bundled `swebench` environment: + +```bash +pip install aenvironment +aenv config init +aenv config set hub_backend http://envhub.aenv.svc.cluster.local:8083/ +export AENV_SYSTEM_URL=http://api-service.aenv.svc.cluster.local:8080 + +cd AEnvironment/aenv/builtin-envs/swebench # config.json: name=swebench, version=1.0.0 +aenv build --push # requires a configured image registry +aenv push # publish metadata to EnvHub +aenv list # `swebench` should now appear +``` + +Its reward entrypoint is `swebench_reward(instance, model_patch, timeout)`, which is +exactly the `call_reward({"instance", "model_patch", "timeout"})` contract +AReaL-SWEAgent uses to grade patches. + +> **Version note:** the repo ships `swebench@1.0.0`, while AReaL-SWEAgent defaults to +> `swebench@1.0.4`. Either export `SWEBENCH_EVAL_ENV=swebench@1.0.0` in the worker +> `env_vars`, or bump the `version` in `config.json` before `aenv build`. + +### 6.3 Provide the task sandbox + +The `persistent-bash-env` / `cc-bash-env` task sandboxes are **not** shipped in the +AEnvironment repo. Use the image already registered in your AEnvironment cluster, or +build your own environment that exposes a shell (the repo's +`aenv/builtin-envs/mini-terminal` is the closest reference) and register it the same way +as section 6.2. AReaL-SWEAgent selects it through the agent config's `aenv_version` +field. + +## 7. Train the Claude Code (cc) agent + +Section 5 runs the built-in `swe` agent. To train the **Claude Code** agent instead — +Claude Code solves the task in a `cc-bash-env` sandbox and the same `swebench` +environment grades the resulting patch — switch `econfig.agent_type` to `cc` and select +a cc config: + +```yaml +econfig: + agent_type: cc + cc_agent_config: train_cc_time3600 # a config under AReaL-SWEAgent/aweagent/configs/ + agent_root: /path/to/AReaL-SWEAgent + timeout: 3600.0 +``` + +The cc agent's LLM calls are still routed through AReaL's OpenAI-compatible proxy (AReaL +injects the per-rollout `base_url`/`api_key`), so training stays on-policy exactly like +the swe path. + +### 7.1 Environment variables for cc + +In addition to the section 2 variables (`AWEAGENT_ROOT`, `PYTHONPATH`, +`AENV_SYSTEM_URL`), the cc agent reads these from the worker `env_vars` — they are +consumed by the cc config in AReaL-SWEAgent (`aweagent/configs/1_0_0/cc.yaml`), not by +AReaL itself: + +| Variable | Purpose | +| ---------------------------------------- | ---------------------------------------------------------- | +| `SWEBENCH_EVAL_ENV` | Grading environment name, e.g. `swebench@1.0.0` (see 6.2). | +| `REMOTE_PROXY_SERVICE_URL` | Gateway that fronts the Claude CLI traffic (proxy mode). | +| `REMOTE_PROXY_API_KEY` | API key for that proxy gateway. | +| `ANTHROPIC_BASE_URL` / `ANTHROPIC_MODEL` | Used instead when running the CLI in direct mode. | + +The concrete agent behaviour — system prompt, `cli_flags`, tool allow-list, +thinking-token budgets — lives entirely in the AReaL-SWEAgent cc config; see that +repository for the available `cc_agent_config` values and their fields. diff --git a/examples/swe/agent.py b/examples/swe/agent.py new file mode 100644 index 0000000000..40b0b8f2aa --- /dev/null +++ b/examples/swe/agent.py @@ -0,0 +1,260 @@ +"""AReaL-SWEAgent workflow for AReaL proxy mode. + +This module runs AReaL-SWEAgent-backed SWE-bench agents through AReaL's proxy +server during RL training. AReaL-SWEAgent stays in a separate checkout and is +discovered through ``AWEAGENT_ROOT`` or ``SWE_AGENT_ROOT``. AReaL passes the +proxy ``base_url`` and per-rollout ``api_key`` explicitly as kwargs instead +of mutating process-wide ``os.environ`` - the latter races across concurrent +rollouts and routes requests to the wrong proxy session. +""" + +import asyncio +import os +import sys +import time +from typing import Any + +from areal.utils import logging + +logger = logging.getLogger("AReaL-SWEAgent") + +_DEFAULT_AGENT_CONFIGS = { + "swe": "1_0_0/min-swe-agent-train-top1", + "cc": "train_cc_time3600", + "oh": "eval_oh", + "opencode": "eval_opencode", + "codex": "eval_codex", +} + + +def _default_aweagent_root() -> str: + """Return the default sibling checkout path for AReaL-SWEAgent.""" + areal_root = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + ) + return os.path.join(os.path.dirname(areal_root), "AReaL-SWEAgent") + + +def _ensure_aweagent_importable(aweagent_root: str = "") -> str: + """Add the AReaL-SWEAgent checkout root to ``sys.path``. + + Resolution order: + 1. Explicit ``aweagent_root`` argument (``econfig.agent_root`` / + ``econfig.aweagent_root`` / ``econfig.swe_agent_root``). + 2. ``AWEAGENT_ROOT`` / ``SWE_AGENT_ROOT`` / ``SWEAGENT_ROOT`` env vars. + 3. ``../AReaL-SWEAgent`` relative to the AReaL repository root. + + Returns the resolved AReaL-SWEAgent checkout root. + """ + if aweagent_root: + root = os.path.abspath(aweagent_root) + elif ( + os.getenv("AWEAGENT_ROOT") + or os.getenv("SWE_AGENT_ROOT") + or os.getenv("SWEAGENT_ROOT") + ): + root = os.path.abspath( + os.getenv("AWEAGENT_ROOT") + or os.getenv("SWE_AGENT_ROOT") + or os.getenv("SWEAGENT_ROOT") + ) + else: + root = _default_aweagent_root() + + if os.path.isdir(root) and root not in sys.path: + sys.path.insert(0, root) + logger.info(f"Added AReaL-SWEAgent root to sys.path: {root}") + elif not os.path.isdir(root): + logger.warning( + f"AReaL-SWEAgent root does not exist: {root}. " + "Set AWEAGENT_ROOT (or SWE_AGENT_ROOT) env var or econfig.agent_root " + "(or the legacy econfig.swe_agent_root) to the correct path." + ) + return root + + +def _configured_aweagent_root(econfig: dict[str, Any]) -> str: + return ( + econfig.get("agent_root") + or econfig.get("aweagent_root") + or econfig.get("swe_agent_root") + or "" + ) + + +def _configured_agent_type(econfig: dict[str, Any]) -> str: + agent_type = str(econfig.get("agent_type") or "swe").strip().lower() + if not agent_type: + raise ValueError("econfig.agent_type must not be empty") + return agent_type + + +def _configured_agent_config(econfig: dict[str, Any], agent_type: str) -> str: + return ( + econfig.get("agent_config") + or econfig.get(f"{agent_type}_agent_config") + or econfig.get("swe_agent_config") + or _DEFAULT_AGENT_CONFIGS.get(agent_type) + or agent_type + ) + + +class SWEAgentWorkflow: + """AReaL-SWEAgent workflow for AReaL proxy mode. + + This workflow runs an AReaL-SWEAgent agent type (``swe``, ``cc``, ``oh``, + ``opencode``, or ``codex``) on SWE-bench style records in sandboxed + environments. The agent's LLM calls are routed through AReaL's proxy + server for on-policy RL training. + + The workflow delegates to AReaL-SWEAgent's ``run_agent_with_reward``, passing + the AReaL proxy URL + per-rollout API key as explicit kwargs so that the + current rollout's endpoint/key are used without process-wide env races. + + Args: + econfig: Environment configuration dict. Key fields include + ``agent_type``, ``agent_config``, ``swe_agent_config``, + ``cc_agent_config``, ``agent_root`` / ``swe_agent_root``, + ``llm_model``, ``opencode_provider``, ``codex_provider``, + and ``timeout``. + gen_args: Generation arguments (unused; token limits come from + the AReaL-SWEAgent config YAML). + timeout: Maximum time allowed for a single episode (default: 1800s). + """ + + def __init__( + self, + econfig: dict | None = None, + gen_args: dict | None = None, + timeout: float = 1800.0, + ): + if econfig is None: + econfig = {} + self.econfig = econfig + self.gen_args = gen_args or {} + self.timeout = econfig.get("timeout", timeout) + + # Ensure AReaL-SWEAgent is importable at construction time so we get an early + # warning if the path is wrong, rather than failing mid-training. + _ensure_aweagent_importable(_configured_aweagent_root(econfig)) + + async def run( + self, data: dict[str, Any], **extra_kwargs: Any + ) -> dict[str, float] | float: + """Run one SWE-bench style agent episode. + + Args: + data: Input data containing SWE-bench instance fields: + - instance_id (str): SWE-bench instance ID + - problem_statement (str): The GitHub issue description + - eval_script (str): Shell script for reward evaluation + **extra_kwargs: Additional kwargs injected by AReaL proxy infra: + - base_url (str): Proxy server URL for the agent LLM. + - api_key (str): Per-rollout API key for the proxy session. + + Returns: + float: The reward from the episode (0.0 or 1.0). + """ + base_url: str | None = extra_kwargs.get("base_url", None) + if base_url is None: + raise ValueError("base_url is required for SWEAgentWorkflow") + + api_key: str = extra_kwargs.get("api_key", os.getenv("OPENAI_API_KEY", "dummy")) + + econfig = self.econfig.copy() + if "econfig" in data: + econfig.update(data["econfig"]) + + agent_type = _configured_agent_type(econfig) + config_name = _configured_agent_config(econfig, agent_type) + llm_model = econfig.get("llm_model") or os.getenv("LLM_MODEL") or None + opencode_provider = ( + econfig.get("opencode_provider") or os.getenv("OPENCODE_PROVIDER") or None + ) + codex_provider = ( + econfig.get("codex_provider") or os.getenv("CODEX_PROVIDER") or None + ) + instance_id = data.get("instance_id", "unknown") + + logger.info( + f"Starting {agent_type} episode: " + f"instance_id={instance_id}, config={config_name}" + ) + start_time = time.time() + + try: + reward = await asyncio.wait_for( + self._run_episode( + data=data, + agent_type=agent_type, + config_name=config_name, + base_url=base_url, + api_key=api_key, + llm_model=llm_model, + opencode_provider=opencode_provider, + codex_provider=codex_provider, + ), + timeout=self.timeout, + ) + except TimeoutError: + elapsed = time.time() - start_time + logger.error( + f"TIMEOUT: Instance {instance_id} exceeded {self.timeout}s " + f"(elapsed: {elapsed:.1f}s). Discarding trajectory." + ) + raise + + elapsed = time.time() - start_time + logger.info( + f"Finished {agent_type} episode: instance_id={instance_id}, " + f"reward={reward}, elapsed={elapsed:.1f}s" + ) + return float(reward) + + async def _run_episode( + self, + data: dict[str, Any], + agent_type: str, + config_name: str, + base_url: str, + api_key: str, + llm_model: str | None, + opencode_provider: str | None, + codex_provider: str | None, + ) -> float: + """Execute one episode through AReaL-SWEAgent and return the reward. + + Delegates to ``aweagent.lifecycle.run_agent_with_reward``, which + handles env creation, the agent loop, reward computation, and trace + persistence. ``base_url`` / ``api_key`` are forwarded as explicit + kwargs (NOT via ``os.environ``) so concurrent rollouts can't read a + foreign session's key. + """ + # Imported lazily - AReaL-SWEAgent is only added to sys.path at run time. + from aweagent.lifecycle import ( + run_agent_with_reward, # type: ignore[import-not-found] + ) + + result_dir = os.getenv("LOG_DIR", "./logs") + try: + reward, _ = await run_agent_with_reward( + data, + agent_type=agent_type, + agent_config=config_name, + result_dir=result_dir, + override_api_key=api_key, + override_base_url=base_url, + override_llm_model=llm_model, + override_opencode_provider=opencode_provider, + override_codex_provider=codex_provider, + ) + except Exception as e: + import traceback + + instance_id = data.get("instance_id", "unknown") + logger.error( + f"[{instance_id}] Episode error: {e}\n{traceback.format_exc()}" + ) + return 0.0 + + return float(reward) diff --git a/examples/swe/filter_function.py b/examples/swe/filter_function.py new file mode 100644 index 0000000000..be397693a0 --- /dev/null +++ b/examples/swe/filter_function.py @@ -0,0 +1,26 @@ +from areal.utils import stats_tracker + + +def filter_function(sample): + # Use original_rewards if available for filtering and statistics + rewards = ( + sample["original_rewards"] + if "original_rewards" in sample + else sample["rewards"] + ) + reward0 = rewards[0] + accept = any(r != reward0 for r in rewards) + + # Track rejection statistics + if not accept: + stats_tracker.get("rollout").scalar(rejected_by_failed_or_perfect=1) + # Distinguish all-correct vs all-wrong + all_positive = all(r > 0 for r in rewards) + stats_tracker.get("rollout").scalar(rejected_by_all_correct=int(all_positive)) + stats_tracker.get("rollout").scalar(rejected_by_all_wrong=int(not all_positive)) + else: + stats_tracker.get("rollout").scalar(rejected_by_failed_or_perfect=0) + stats_tracker.get("rollout").scalar(rejected_by_all_correct=0) + stats_tracker.get("rollout").scalar(rejected_by_all_wrong=0) + + return accept diff --git a/examples/swe/prefix_matchers.py b/examples/swe/prefix_matchers.py new file mode 100644 index 0000000000..64b52bdf79 --- /dev/null +++ b/examples/swe/prefix_matchers.py @@ -0,0 +1,58 @@ +"""Custom prefix matchers for InteractionCache parent-child matching. + +The default cache matcher requires exact message equality. SWE-bench agents, +especially Claude Code style agents routed through Anthropic-compatible APIs, +can rewrite tool arguments or tool output formatting between turns while still +representing the same conversation. The matcher here keeps concat export stable +for those known rewrites. +""" + +from __future__ import annotations + + +def _tool_call_ids(tool_calls: list[dict]) -> list[str]: + """Extract ordered tool_call IDs from a tool_calls list.""" + return [tc.get("id", "") for tc in tool_calls if isinstance(tc, dict)] + + +def _messages_match(a: dict, b: dict) -> bool: + """Check whether two message dicts are semantically equivalent.""" + if a.get("role") != b.get("role"): + return False + + role = a.get("role") + + if role == "tool": + return a.get("tool_call_id") == b.get("tool_call_id") + + if a.get("content", "") != b.get("content", ""): + return False + + if a.get("thinking") != b.get("thinking"): + return False + + a_tc = a.get("tool_calls") + b_tc = b.get("tool_calls") + if a_tc is not None or b_tc is not None: + if a_tc is None or b_tc is None: + return False + if not isinstance(a_tc, list) or not isinstance(b_tc, list): + return a_tc == b_tc + if len(a_tc) != len(b_tc): + return False + if _tool_call_ids(a_tc) != _tool_call_ids(b_tc): + return False + + return True + + +def swe_prefix_matcher(a: list[dict], b: list[dict]) -> bool: + """Return True if ``a`` is a semantic prefix of ``b``.""" + if len(a) > len(b): + return False + for am, bm in zip(a, b): + if am == bm: + continue + if not _messages_match(am, bm): + return False + return True diff --git a/examples/swe/preprocessors.py b/examples/swe/preprocessors.py new file mode 100644 index 0000000000..8fff09ccee --- /dev/null +++ b/examples/swe/preprocessors.py @@ -0,0 +1,124 @@ +"""Message preprocessors for SWE-bench agent rollouts. + +Preprocessors transform messages after Anthropic-to-OpenAI translation and +before they reach AReaL's OpenAI client. They remove volatile Claude Code +metadata that otherwise breaks prefix-based parent matching in concat export. +""" + +import json +import re +from typing import Protocol, runtime_checkable + + +@runtime_checkable +class MessagePreprocessor(Protocol): + """Protocol for message preprocessors.""" + + def __call__(self, messages: list[dict]) -> list[dict]: ... + + +class StripAnthropicBillingHeader: + """Remove per-request Anthropic billing header lines from system prompts.""" + + _PATTERN = re.compile(r"^x-anthropic-billing-header:[^\n]*\n?", re.MULTILINE) + + def __call__(self, messages: list[dict]) -> list[dict]: + if not messages: + return messages + first = messages[0] + if first.get("role") == "system" and isinstance(first.get("content"), str): + first["content"] = self._PATTERN.sub("", first["content"]) + return messages + + +class NormalizeSystemReminder: + """Remove volatile ``currentDate`` lines from system reminders.""" + + _PATTERN = re.compile( + r"# currentDate\nToday's date is \d{4}-\d{2}-\d{2}\.\n?", + ) + + def __call__(self, messages: list[dict]) -> list[dict]: + for msg in messages: + content = msg.get("content") + if isinstance(content, str) and "currentDate" in content: + msg["content"] = self._PATTERN.sub("", content) + return messages + + +class StripAllSystemReminders: + """Remove all ``...`` blocks.""" + + _PATTERN = re.compile(r"[\s\S]*?") + + def __call__(self, messages: list[dict]) -> list[dict]: + for msg in messages: + content = msg.get("content") + if isinstance(content, str) and "" in content: + msg["content"] = self._PATTERN.sub("", content) + return messages + + +class StripAnthropicCacheFields: + """Strip Anthropic-specific fields that are not preserved in stored output.""" + + def __call__(self, messages: list[dict]) -> list[dict]: + for msg in messages: + msg.pop("cache_control", None) + msg.pop("thinking_blocks", None) + if "tool_calls" in msg: + for tc in msg["tool_calls"]: + if isinstance(tc, dict): + tc.pop("cache_control", None) + tc.pop("provider_specific_fields", None) + fn = tc.get("function") + if isinstance(fn, dict): + fn.pop("cache_control", None) + return messages + + +class NormalizeToolCallArguments: + """Normalize tool_call arguments for deterministic comparison.""" + + _TOOL_ARG_DEFAULTS: dict[tuple[str, str], object] = { + ("Edit", "replace_all"): False, + } + + @classmethod + def _normalize_tool_arguments( + cls, + tool_name: str | None, + args_dict: dict, + ) -> dict: + if tool_name is not None: + for (tn, field), default in cls._TOOL_ARG_DEFAULTS.items(): + if tn == tool_name and args_dict.get(field) == default: + args_dict.pop(field, None) + return args_dict + + def __call__(self, messages: list[dict]) -> list[dict]: + for msg in messages: + if "tool_calls" in msg: + for tc in msg["tool_calls"]: + if isinstance(tc, dict): + fn = tc.get("function") + if isinstance(fn, dict): + tool_name = fn.get("name") + args = fn.get("arguments") + if isinstance(args, str): + try: + parsed = json.loads(args) + parsed = self._normalize_tool_arguments( + tool_name, + parsed, + ) + fn["arguments"] = json.dumps( + parsed, + sort_keys=True, + ensure_ascii=False, + ) + except (json.JSONDecodeError, TypeError): + pass + if msg.get("role") == "assistant" and "content" not in msg: + msg["content"] = "" + return messages diff --git a/examples/swe/qwen3_30b_a3b_grpo.yaml b/examples/swe/qwen3_30b_a3b_grpo.yaml new file mode 100644 index 0000000000..a9fc7c3c34 --- /dev/null +++ b/examples/swe/qwen3_30b_a3b_grpo.yaml @@ -0,0 +1,291 @@ +experiment_name: qwen3_30b_a3b_grpo +trial_name: trial0 + +seed: 42 +enable_offload: true +total_train_epochs: 10 +total_train_steps: 100 +tokenizer_path: ${actor.path} + +# Clean up external SWE sandbox instances after the run (LOG_DIR is injected). +post_exit_hook: 'AWE_ROOT=\${AWEAGENT_ROOT:-\${SWE_AGENT_ROOT:-/path/to/AReaL-SWEAgent}}; cd "$AWE_ROOT" && PYTHONPATH=$PWD:$PYTHONPATH pip install -q -r requirements.txt && AENV_SYSTEM_URL=\${AENV_SYSTEM_URL} python -m aweagent.maintenance.clean_instances' + +cluster: + # 4 nodes x 8 GPUs: 2 for rollout (SGLang), 2 for actor (Megatron). + n_nodes: 4 + n_gpus_per_node: 8 + fileroot: /tmp/areal/experiments + name_resolve: + type: nfs + nfs_record_root: /tmp/areal/name_resolve + +scheduler: + type: slurm + +gconfig: + n_samples: 4 + min_new_tokens: 0 + max_new_tokens: 131071 + max_tokens: 131071 + greedy: false + temperature: 1.0 + top_p: 1.0 + top_k: 1000000 + +rollout: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + backend: "sglang:d4t8p1" + max_concurrent_rollouts: 32 + queue_size: null + consumer_batch_size: ${train_dataset.batch_size} + max_head_offpolicyness: 4 + enable_rollout_tracing: true + scheduling_spec: ${actor.scheduling_spec} + fileroot: ${cluster.fileroot} + tokenizer_path: ${tokenizer_path} + dump_to_file: true + setup_timeout: 900.0 + request_timeout: 20000.0 + request_retries: 3 + pause_grace_period: 30 + agent: + mode: inline + # admin_api_key is injected by the launcher; not stored in YAML. + admin_api_key: ${oc.env:SWE_RL_ADMIN_API_KEY} + tool_call_parser: qwen3_coder + reasoning_parser: qwen3 + chat_template_type: concat + engine_max_tokens: ${gconfig.max_tokens} + export_style: concat + turn_discount: 1.0 + session_timeout_seconds: 3600 + prefix_matcher: examples.swe.prefix_matchers.swe_prefix_matcher + message_preprocessors: + - examples.swe.preprocessors.StripAnthropicBillingHeader + - examples.swe.preprocessors.StripAllSystemReminders + - examples.swe.preprocessors.StripAnthropicCacheFields + - examples.swe.preprocessors.NormalizeToolCallArguments + +actor: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + # attn DP4 PP1 TP4 CP2; ffn DP4 PP1 EP8. CP=2 shards long sequences. + backend: "megatron:(attn:d4p1t4c2|ffn:d4p1e8)" + path: /path/to/Qwen3-Coder-30B-A3B-Instruct + init_from_scratch: false + disable_dropout: true + gradient_checkpointing: true + dtype: bfloat16 + grad_reduce_dtype: float32 + mb_spec: + n_mbs: 1 + granularity: 1 + max_tokens_per_mb: 131072 + n_mbs_divisor: 1 + optimizer: + type: adam + lr: 3.0e-6 + weight_decay: 0.01 + beta1: 0.9 + beta2: 0.999 + eps: 1.0e-8 + lr_scheduler_type: constant + gradient_clipping: 1.0 + warmup_steps_proportion: 0.001 + megatron: + wrap_with_ddp: true + ddp: + grad_reduce_in_fp32: true + overlap_grad_reduce: true + overlap_param_gather: false + use_distributed_optimizer: true + use_deterministic_algorithms: true + enable_fp32_lm_head: false + cross_entropy_loss_fusion: false + recompute_granularity: full + recompute_method: uniform + recompute_num_layers: 1 + moe_router_dtype: fp32 + moe_token_dispatcher_type: alltoall + moe_router_bias_update_rate: 0.0 + eps_clip: 0.2 + eps_clip_higher: 0.28 + temperature: ${gconfig.temperature} + kl_ctl: 0.0 + ppo_n_minibatches: 1 + recompute_logprob: false + prox_logp_method: recompute + use_decoupled_loss: true + rejection_sampling: + level: token + action: mask + metric: ratio + upper: 5.0 + discount: 1.0 + gae_lambda: 1.0 + reward_scaling: 1.0 + reward_bias: 0.0 + reward_clip: 5.0 + reward_norm: + mean_level: group + std_level: group + group_size: ${gconfig.n_samples} + eps: 1.0e-5 + adv_norm: + mean_level: null + std_level: null + max_new_tokens: ${gconfig.max_new_tokens} + scheduling_spec: + - task_type: worker + port_count: 2 + gpu: 1 + cpu: 4 + mem: 32 + exclusive: true + reservation: null + nodelist: null + exclude: null + image: /path/to/areal.sif + cmd: python3 -m areal.infra.rpc.rpc_server + env_vars: + NCCL_TIMEOUT: 600 + PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True" + # Point HOME and caches at shared storage to avoid filling the + # container's small default HOME during JIT compilation. + HOME: "/path/to/cache/home" + HF_HOME: "/path/to/cache/hf" + XDG_CACHE_HOME: "/path/to/cache/xdg" + # Workers must import both AReaL and the external AReaL-SWEAgent. + PYTHONPATH: "/path/to/AReaL-SWEAgent:/path/to/AReaL" + AENV_SYSTEM_URL: "http://your-aenv-service:8080" + SWE_AGENT_ROOT: "/path/to/AReaL-SWEAgent" + AWEAGENT_ROOT: "/path/to/AReaL-SWEAgent" + OPENAI_MODEL: "Qwen3-Coder-30B-A3B-Instruct" + additional_bash_cmds: + # Install AReaL-SWEAgent's aenv/fastmcp into the venv the workers use. + - "/opt/.venv/bin/python -m ensurepip --upgrade >/dev/null 2>&1 || true" + - "/opt/.venv/bin/python -m pip install --no-deps aenvironment==0.1.8rc2 || echo 'WARN: aenvironment install failed'" + - "/opt/.venv/bin/python -m pip install 'fastmcp<3' || echo 'WARN: fastmcp install failed'" + - "/opt/.venv/bin/python -c 'import aenv, aweagent' || echo 'WARN: aenv/aweagent import failed'" + +ref: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + backend: ${actor.backend} + path: ${actor.path} + init_from_scratch: false + disable_dropout: false + dtype: ${actor.dtype} + mb_spec: + n_mbs: 1 + granularity: 1 + max_tokens_per_mb: 131072 + n_mbs_divisor: 1 + optimizer: null + scheduling_strategy: + type: colocation + target: actor + scheduling_spec: ${actor.scheduling_spec} + +# SWE-bench agent environment. Set agent_type=cc to train Claude Code through +# the same proxy path. agent_root points at the external AReaL-SWEAgent checkout. +econfig: + dataset_path: "" + agent_type: swe + agent_config: "" + swe_agent_config: 1_0_0/min-swe-agent-train-top1 + cc_agent_config: train_cc_time3600 + agent_root: /path/to/AReaL-SWEAgent + swe_agent_root: /path/to/AReaL-SWEAgent + llm_model: "" + opencode_provider: "" + codex_provider: "" + max_completion_tokens: 131071 + timeout: 3600.0 + +sglang: + model_path: ${actor.path} + random_seed: ${seed} + skip_tokenizer_init: true + dtype: ${actor.dtype} + context_length: 131072 + max_prefill_tokens: 131071 + max_running_requests: 16 + mem_fraction_static: 0.85 + chunked_prefill_size: 8192 + cuda_graph_max_bs: 16 + schedule_policy: fcfs + attention_backend: fa3 + log_level: info + log_requests: false + sampling_backend: pytorch + enable_memory_saver: false + disable_radix_cache: false + triton_attention_num_kv_splits: 16 + disable_custom_all_reduce: true + disable_cuda_graph: false + +# Datasets: SWE-bench style jsonl with instance_id / problem_statement / +# eval_script per line. +train_dataset: + batch_size: 4 + shuffle: true + pin_memory: true + num_workers: 4 + path: /path/to/swe_bench_rl.jsonl + type: rl + max_length: 32768 + drop_last: true + +valid_dataset: + batch_size: 4 + shuffle: true + pin_memory: true + num_workers: 4 + path: /path/to/swe_bench_rl.jsonl + type: rl + +# Utilities +saver: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: 10 + freq_secs: null + +recover: + mode: auto + retries: 3 + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_secs: 7200 + no_save_optim: false + no_load_optim: false + freq_steps: 10 + +evaluator: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: null + freq_steps: null + freq_secs: null + +stats_logger: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + wandb: + mode: disabled + +perf_tracer: + enabled: false + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + session_tracer: + enabled: false diff --git a/examples/swe/train_swe_rl.py b/examples/swe/train_swe_rl.py new file mode 100644 index 0000000000..57db957ef7 --- /dev/null +++ b/examples/swe/train_swe_rl.py @@ -0,0 +1,218 @@ +"""Training script for SWE-bench agent RL with AReaL proxy mode.""" + +import json +import sys +import warnings +from pathlib import Path +from typing import Any + +from datasets import Dataset + +from examples.swe.utils import SWEPPOConfig + +from areal import PPOTrainer +from areal.api.cli_args import load_expr_config +from areal.utils import logging + +logger = logging.getLogger("SWETrain") + + +def get_swe_dataset( + dataset_path: str, + split: str = "train", + min_items: int = 64, +) -> Dataset: + """Create a HuggingFace Dataset from a SWE-bench JSONL file. + + Each line in the JSONL file should be a SWE-bench instance with at minimum: + - instance_id: The SWE-bench instance ID (e.g., "django__django-10097") + - problem_statement: The GitHub issue description + - eval_script: Shell script to evaluate the agent's fix + + Args: + dataset_path: Path to the SWE-bench JSONL file. + split: Informational split label (not used for filtering). + min_items: Minimum dataset size; items are duplicated if fewer exist. + + Returns: + HuggingFace Dataset of SWE-bench instances. + """ + path = Path(dataset_path) + if not path.exists(): + raise FileNotFoundError(f"SWE-bench dataset not found: {dataset_path}") + + dataset_items = [] + with open(path, encoding="utf-8") as f: + for line in f: + line = line.strip() + if not line: + continue + item = json.loads(line) + if "instance_id" not in item: + logger.warning(f"Skipping item missing 'instance_id': {line[:100]}") + continue + if "problem_statement" not in item: + logger.warning( + "Skipping item missing 'problem_statement': " + f"{item.get('instance_id')}" + ) + continue + dataset_items.append(item) + + if not dataset_items: + raise ValueError(f"No valid items found in dataset: {dataset_path}") + + # Duplicate dataset if fewer than min_items for efficient batching + if len(dataset_items) < min_items: + original_items = dataset_items.copy() + while len(dataset_items) < min_items: + dataset_items.extend(original_items) + + dataset = Dataset.from_list(dataset_items) + logger.info( + f"Created SWE dataset with {len(dataset)} items " + f"from {dataset_path} (split={split})" + ) + return dataset + + +def group_filter(x: dict[str, Any]): + """Filter out groups where all rollouts already solved the task.""" + return x["rewards"].mean() <= 0.95 + + +def _install_aweagent_deps_on_ray_nodes(aweagent_root: str): + """Install AReaL-SWEAgent dependencies on all Ray GPU nodes. + + Each node runs in a separate container with its own venv, + so we must ensure packages like ``aenv`` are installed everywhere. + """ + if not aweagent_root: + aweagent_root = str( + Path(__file__).resolve().parents[2].parent / "AReaL-SWEAgent" + ) + try: + import ray + + if not ray.is_initialized(): + return + + @ray.remote(num_gpus=0) + def _install(): + import os + import socket + import subprocess + + ip = socket.gethostbyname(socket.gethostname()) + req_path = os.path.join(aweagent_root, "requirements.txt") + result = subprocess.run( + ["uv", "pip", "install", "-r", req_path], + capture_output=True, + text=True, + timeout=120, + ) + return ( + ip, + result.returncode, + result.stderr[-200:] if result.stderr else "", + ) + + nodes = [ + n + for n in ray.nodes() + if n.get("Alive") and n.get("Resources", {}).get("GPU", 0) > 0 + ] + refs = [] + for node in nodes: + node_ip = node["NodeManagerAddress"] + refs.append(_install.options(resources={f"node:{node_ip}": 0.01}).remote()) + + results = ray.get(refs, timeout=180) + for ip, rc, err in results: + if rc != 0: + logger.warning(f"Failed to install AReaL-SWEAgent deps on {ip}: {err}") + else: + logger.info(f"AReaL-SWEAgent deps installed on {ip}") + except Exception as e: + logger.warning(f"Could not install AReaL-SWEAgent deps on Ray nodes: {e}") + + +def _resolve_aweagent_root(econfig) -> str: + return ( + getattr(econfig, "agent_root", "") + or getattr(econfig, "aweagent_root", "") + or getattr(econfig, "swe_agent_root", "") + ) + + +def main(args): + warnings.filterwarnings("ignore", category=UserWarning, module="pydantic") + + config, _ = load_expr_config(args, SWEPPOConfig) + econfig = config.econfig + + # When using Ray scheduler, ensure SWEAgent deps are on all nodes + if config.scheduler.type == "ray": + import ray + + ray.init(address="auto", ignore_reinit_error=True) + _install_aweagent_deps_on_ray_nodes(_resolve_aweagent_root(econfig)) + + # Resolve dataset paths from config + train_path = config.train_dataset.path + valid_path = config.valid_dataset.path + + def resolve_path(p: str) -> str: + if Path(p).is_absolute() or Path(p).exists(): + return p + if econfig.dataset_path: + candidate = Path(econfig.dataset_path) / p + if candidate.exists(): + return str(candidate) + return p + + train_dataset = get_swe_dataset( + dataset_path=resolve_path(train_path), + split="train", + ) + valid_dataset = get_swe_dataset( + dataset_path=resolve_path(valid_path), + split="test", + ) + + # Build workflow kwargs + from dataclasses import asdict + + econfig_dict = asdict(econfig) + workflow_kwargs = dict( + econfig=econfig_dict, + gen_args=dict( + temperature=config.gconfig.temperature, + max_completion_tokens=config.gconfig.max_new_tokens, + ), + timeout=econfig.timeout, + ) + + # Eval workflow with lower temperature for deterministic evaluation + eval_workflow_kwargs = workflow_kwargs.copy() + eval_workflow_kwargs["gen_args"] = dict( + temperature=0.0, + max_completion_tokens=config.gconfig.max_new_tokens, + ) + + with PPOTrainer( + config, + train_dataset=train_dataset, + valid_dataset=valid_dataset, + ) as trainer: + trainer.train( + workflow="examples.swe.agent.SWEAgentWorkflow", + workflow_kwargs=workflow_kwargs, + eval_workflow=None, + eval_workflow_kwargs=eval_workflow_kwargs, + dynamic_filter_fn=getattr(config, "should_accept_fn", None), + ) + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/examples/swe/utils.py b/examples/swe/utils.py new file mode 100644 index 0000000000..c53881cbff --- /dev/null +++ b/examples/swe/utils.py @@ -0,0 +1,123 @@ +"""Utilities for SWE-bench agent training with AReaL.""" + +from dataclasses import dataclass, field + +from areal.api.cli_args import PPOConfig + + +@dataclass +class SWEEnvConfig: + """Environment configuration for AReaL-SWEAgent-backed SWE-bench training. + + Attributes: + dataset_path: Path to the SWE-bench JSONL dataset file. + agent_type: AReaL-SWEAgent agent type to train, e.g. ``swe`` or ``cc``. + agent_config: Generic AReaL-SWEAgent config name. When set, this overrides + the compatibility fields below. + swe_agent_config: Compatibility config field for ``agent_type=swe``. + cc_agent_config: Compatibility config field for ``agent_type=cc``. + agent_root: Root directory of the external AReaL-SWEAgent checkout. + swe_agent_root: Legacy alias for ``agent_root``. + llm_model: Optional LLM model override for OH/OpenCode/Codex agents. + opencode_provider: Optional OpenCode provider override. + codex_provider: Optional Codex provider override. + step_limit: Maximum number of agent interaction steps per episode. + max_completion_tokens: Maximum completion tokens for the agent LLM. + timeout: Maximum time allowed for a single episode in seconds. + """ + + dataset_path: str = field( + default="", + metadata={"help": "Path to the SWE-bench JSONL dataset file."}, + ) + agent_type: str = field( + default="swe", + metadata={ + "help": ( + "AReaL-SWEAgent agent type to run. Supported by AReaL-SWEAgent main: " + "'swe', 'cc', 'oh', 'opencode', and 'codex'." + ) + }, + ) + agent_config: str = field( + default="", + metadata={ + "help": ( + "Generic AReaL-SWEAgent YAML config name. When non-empty, overrides " + "swe_agent_config / cc_agent_config." + ) + }, + ) + swe_agent_config: str = field( + default="1_0_0/min-swe-agent-train-top1", + metadata={ + "help": ( + "Name of the AReaL-SWEAgent YAML config under the external AReaL-SWEAgent " + "checkout. Defaults to the Qwen SWE-RL training config." + ) + }, + ) + cc_agent_config: str = field( + default="train_cc_time3600", + metadata={ + "help": ( + "Name of the AReaL-SWEAgent YAML config used when agent_type='cc'. " + "Kept separate for compatibility with swe/main configs." + ) + }, + ) + agent_root: str = field( + default="", + metadata={ + "help": ( + "Root directory of the external AReaL-SWEAgent checkout. Defaults to " + "../AReaL-SWEAgent relative to the AReaL repository when unset." + ) + }, + ) + swe_agent_root: str = field( + default="", + metadata={ + "help": ( + "Legacy alias for agent_root / AWEAGENT_ROOT. Kept so older " + "SWE launch scripts keep working." + ) + }, + ) + llm_model: str = field( + default="", + metadata={"help": "Optional model name override for OH/OpenCode/Codex agents."}, + ) + opencode_provider: str = field( + default="", + metadata={"help": "Optional provider override for OpenCode agents."}, + ) + codex_provider: str = field( + default="", + metadata={"help": "Optional provider override for Codex agents."}, + ) + step_limit: int = field( + default=100, + metadata={"help": "Maximum number of agent interaction steps per episode."}, + ) + max_completion_tokens: int = field( + default=16384, + metadata={"help": "Maximum completion tokens for the agent LLM."}, + ) + timeout: float = field( + default=1800.0, + metadata={"help": "Maximum time allowed for a single episode in seconds."}, + ) + + +@dataclass +class SWEPPOConfig(PPOConfig): + """PPO configuration with SWE-bench-specific settings.""" + + econfig: SWEEnvConfig = field(default_factory=SWEEnvConfig) + should_accept_fn: str | None = field( + default=None, + metadata={ + "help": "Import path of the filter function for accepting rollout samples." + }, + ) diff --git a/examples/vlm/clevr_count_70k_grpo.py b/examples/vlm/clevr_count_70k_grpo.py index 8cab44e8ff..2838dc3da4 100644 --- a/examples/vlm/clevr_count_70k_grpo.py +++ b/examples/vlm/clevr_count_70k_grpo.py @@ -4,8 +4,11 @@ from areal import PPOTrainer from areal.api.cli_args import GRPOConfig, load_expr_config from areal.dataset import get_custom_dataset +from areal.utils import logging from areal.utils.hf_utils import load_hf_processor_and_tokenizer +logger = logging.getLogger("clevr_count_70k_grpo") + def extract_answer(pred_str, data_name, use_last_number=True): match = re.findall(r"\[([0-9\.]+)\]", pred_str) @@ -18,15 +21,17 @@ def extract_answer(pred_str, data_name, use_last_number=True): def clevr_count_70k_reward_fn( prompt, completions, prompt_ids, completion_ids, answer, **kwargs ): - sol = extract_answer(completions, data_name="") # str number - ans = answer + try: + sol = extract_answer(str(completions), data_name="") # str number + ans = str(answer) - if sol is None: - return 0 - if ans is None: - return 0 + if not sol or not ans: + return 0.0 - return float(sol.strip() == ans.strip()) + return float(sol.strip() == ans.strip()) + except Exception: + logger.warning("Exception in clevr_count_70k_reward_fn", exc_info=True) + return 0.0 def main(args): diff --git a/examples/vlm/qwen3_6_35b_a3b_megatron_geometry3k_grpo.yaml b/examples/vlm/qwen3_6_35b_a3b_megatron_geometry3k_grpo.yaml new file mode 100644 index 0000000000..f6ceca39b3 --- /dev/null +++ b/examples/vlm/qwen3_6_35b_a3b_megatron_geometry3k_grpo.yaml @@ -0,0 +1,186 @@ +experiment_name: geometry3k-grpo +trial_name: trial1 + +seed: 1 +enable_offload: false +total_train_epochs: 10 +tokenizer_path: ${actor.path} + +cluster: + n_nodes: 6 + n_gpus_per_node: 8 + fileroot: /tmp/areal/experiments + name_resolve: + type: nfs + nfs_record_root: /tmp/areal/name_resolve + +scheduler: + type: slurm + +rollout: + backend: "vllm:d8p1t2" + experiment_name: ${experiment_name} + trial_name: ${trial_name} + max_concurrent_rollouts: 256 + queue_size: null + consumer_batch_size: ${train_dataset.batch_size} + max_head_offpolicyness: 4 + enable_rollout_tracing: false + scheduling_spec: ${actor.scheduling_spec} + fileroot: ${cluster.fileroot} + tokenizer_path: ${tokenizer_path} + dump_to_file: true + +gconfig: + n_samples: 4 + min_new_tokens: 0 + max_new_tokens: 2048 + greedy: false + temperature: 1.0 + +actor: + backend: "megatron:(attn:d8p2t2|ffn:d4p2t1e4)" + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: Qwen/Qwen3.6-35B-A3B + init_from_scratch: false + disable_dropout: true + gradient_checkpointing: true + dtype: bfloat16 + mb_spec: + max_tokens_per_mb: 8192 + optimizer: + type: adam + lr: 1e-6 + weight_decay: 0.01 + beta1: 0.9 + beta2: 0.999 + eps: 1e-8 + lr_scheduler_type: constant + gradient_clipping: 1.0 + warmup_steps_proportion: 0.001 + eps_clip: 0.4 + temperature: ${gconfig.temperature} + reward_scaling: 10.0 + reward_bias: -0.5 + kl_ctl: 0.0 + ppo_n_minibatches: 1 + recompute_logprob: true + use_decoupled_loss: true + rejection_sampling: + metric: ratio + upper: 5.0 + reward_norm: + mean_level: group + std_level: group + group_size: ${gconfig.n_samples} + adv_norm: + mean_level: batch + std_level: batch + megatron: + bridge_type: megatron-bridge + use_bridge_for_update_weights: true + enable_mtp: false + max_new_tokens: ${gconfig.max_new_tokens} + scheduling_spec: + - task_type: worker + port_count: 2 + gpu: 1 + mem: 32 + cmd: python3 -m areal.infra.rpc.rpc_server + env_vars: { PYTORCH_ALLOC_CONF: "expandable_segments:True" } + +ref: + backend: ${actor.backend} + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: ${actor.path} + init_from_scratch: false + disable_dropout: true + dtype: ${actor.dtype} + mb_spec: + max_tokens_per_mb: 4096 + optimizer: null + scheduling_strategy: + type: colocation + target: actor + scheduling_spec: ${actor.scheduling_spec} + +# SGLang +sglang: + model_path: ${actor.path} + random_seed: ${seed} + skip_tokenizer_init: true + dtype: ${actor.dtype} + max_running_requests: 64 + context_length: 32768 + mem_fraction_static: 0.8 + enable_multimodal: true + +vllm: + model: ${actor.path} + seed: ${seed} + skip_tokenizer_init: false + dtype: ${actor.dtype} + max_model_len: 32768 + gpu_memory_utilization: 0.9 + disable_sliding_window: false + enforce_eager: true + gdn_prefill_backend: triton + +# datasets +train_dataset: + batch_size: 128 + shuffle: true + pin_memory: true + num_workers: 4 + path: hiyouga/geometry3k + type: rl + +valid_dataset: + batch_size: 128 + pin_memory: true + num_workers: 4 + path: hiyouga/geometry3k + type: rl + +# Utilities +saver: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +recover: + mode: disabled + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: 3600 + +evaluator: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +stats_logger: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + wandb: + mode: disabled + +perf_tracer: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + enabled: false + session_tracer: + enabled: false diff --git a/pyproject.toml b/pyproject.toml index 1d90e3e717..b325587691 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ description = "AReaL: A Large-Scale Asynchronous Reinforcement Learning System" readme = "README.md" license = {text = "Apache-2.0"} requires-python = ">=3.11,<3.13" -version = "1.0.4" +version = "2.0.0" authors = [ {name = "AReaL Team"}, ] @@ -183,6 +183,13 @@ cuda = [ sandbox = [ "daytona>=0.167.0", ] +cli = [ + "click>=8.1", + "colorlog", +] + +[project.scripts] +areal = "areal.v2.cli.main:cli" [project.urls] "Homepage" = "https://github.com/areal-project/AReaL" diff --git a/pyproject.vllm.toml b/pyproject.vllm.toml index 4ab022c367..c71d4049d0 100644 --- a/pyproject.vllm.toml +++ b/pyproject.vllm.toml @@ -24,7 +24,7 @@ description = "AReaL: A Large-Scale Asynchronous Reinforcement Learning System" readme = "README.md" license = {text = "Apache-2.0"} requires-python = ">=3.11,<3.13" -version = "1.0.4" +version = "2.0.0" authors = [ {name = "AReaL Team"}, ] @@ -194,6 +194,13 @@ cuda = [ sandbox = [ "daytona>=0.167.0", ] +cli = [ + "click>=8.1", + "colorlog", +] + +[project.scripts] +areal = "areal.v2.cli.main:cli" [project.urls] "Homepage" = "https://github.com/areal-project/AReaL" diff --git a/tests/experimental/agent_service/test_integration.py b/tests/experimental/agent_service/test_integration.py deleted file mode 100644 index f815c5ca25..0000000000 --- a/tests/experimental/agent_service/test_integration.py +++ /dev/null @@ -1,288 +0,0 @@ -"""Integration tests for the Agent Service. - -Tests the full HTTP microservice stack: Worker → DataProxy → Router, -plus utility functions from the Bridge and Gateway health endpoints. -""" - -from __future__ import annotations - -from unittest.mock import patch - -import pytest - -from areal.experimental.agent_service.auth import DEFAULT_ADMIN_API_KEY, admin_headers -from areal.experimental.agent_service.data_proxy.app import create_data_proxy_app -from areal.experimental.agent_service.data_proxy.config import DataProxyConfig -from areal.experimental.agent_service.gateway.app import create_gateway_app -from areal.experimental.agent_service.gateway.bridge import OpenResponsesBridge -from areal.experimental.agent_service.gateway.config import GatewayConfig -from areal.experimental.agent_service.router.app import create_router_app -from areal.experimental.agent_service.router.config import RouterConfig -from areal.experimental.agent_service.types import ( - AgentRequest, - AgentResponse, - EventEmitter, -) -from areal.experimental.agent_service.worker.app import create_worker_app - -httpx = pytest.importorskip("httpx") - -_AUTH = admin_headers(DEFAULT_ADMIN_API_KEY) - - -class _EchoAgent: - async def run( - self, request: AgentRequest, *, emitter: EventEmitter - ) -> AgentResponse: - history_summary = f"history={len(request.history)}" - await emitter.emit_delta(f"echo: {request.message} ({history_summary})") - return AgentResponse(summary=f"echo: {request.message}") - - -class _ToolAgent: - async def run( - self, request: AgentRequest, *, emitter: EventEmitter - ) -> AgentResponse: - await emitter.emit_tool_call("lookup", '{"id": "123"}') - await emitter.emit_tool_result("lookup", '{"status": "ok"}') - await emitter.emit_delta("Lookup complete") - return AgentResponse( - summary="Lookup complete", - metadata={"tool_calls": [{"name": "lookup", "arguments": {"id": "123"}}]}, - ) - - -def _make_worker_app(agent_cls): - with patch( - "areal.experimental.agent_service.worker.app.import_from_string", - return_value=agent_cls, - ): - return create_worker_app("mock.path") - - -class TestWorkerDataProxyIntegration: - """Test DataProxy → Worker chain using ASGITransport for the Worker.""" - - @pytest.mark.asyncio - async def test_single_turn(self): - worker_app = _make_worker_app(_EchoAgent) - worker_transport = httpx.ASGITransport(app=worker_app) - - async with httpx.AsyncClient( - transport=worker_transport, base_url="http://worker" - ) as worker_client: - # DataProxy forwards to worker — test worker directly first - resp = await worker_client.post( - "/run", - json={ - "message": "hello", - "session_key": "s1", - "run_id": "r1", - "history": [], - }, - ) - assert resp.status_code == 200 - data = resp.json() - assert "echo: hello" in data["summary"] - - @pytest.mark.asyncio - async def test_data_proxy_manages_history(self): - worker_app = _make_worker_app(_EchoAgent) - worker_transport = httpx.ASGITransport(app=worker_app) - - # Create DataProxy pointing to worker - proxy_app = create_data_proxy_app(DataProxyConfig(worker_addr="http://worker")) - - # Patch DataProxy's httpx client to use worker's ASGITransport - original_post = httpx.AsyncClient.post - - async def patched_post(self, url, **kwargs): - if "worker" in url: - async with httpx.AsyncClient( - transport=worker_transport, base_url="http://worker" - ) as wc: - path = url.split("http://worker")[-1] - return await wc.post(path, **kwargs) - return await original_post(self, url, **kwargs) - - proxy_transport = httpx.ASGITransport(app=proxy_app) - - with patch.object(httpx.AsyncClient, "post", patched_post): - async with httpx.AsyncClient( - transport=proxy_transport, base_url="http://proxy" - ) as proxy_client: - # Turn 1 - r1 = await proxy_client.post( - "/session/s1/turn", - json={"message": "hello", "run_id": "r1"}, - ) - assert r1.status_code == 200 - assert "echo: hello" in r1.json()["summary"] - - # Turn 2 — history should have turn 1 - r2 = await proxy_client.post( - "/session/s1/turn", - json={"message": "world", "run_id": "r2"}, - ) - assert r2.status_code == 200 - - # Check history grew - h = await proxy_client.get("/session/s1/history") - history = h.json()["history"] - assert len(history) >= 2 # at least user+assistant from turn 1 - - @pytest.mark.asyncio - async def test_close_session_clears_history(self): - worker_app = _make_worker_app(_EchoAgent) - worker_transport = httpx.ASGITransport(app=worker_app) - proxy_app = create_data_proxy_app(DataProxyConfig(worker_addr="http://worker")) - - original_post = httpx.AsyncClient.post - - async def patched_post(self, url, **kwargs): - if "worker" in url: - async with httpx.AsyncClient( - transport=worker_transport, base_url="http://worker" - ) as wc: - path = url.split("http://worker")[-1] - return await wc.post(path, **kwargs) - return await original_post(self, url, **kwargs) - - proxy_transport = httpx.ASGITransport(app=proxy_app) - - with patch.object(httpx.AsyncClient, "post", patched_post): - async with httpx.AsyncClient( - transport=proxy_transport, base_url="http://proxy" - ) as proxy_client: - await proxy_client.post( - "/session/s1/turn", - json={"message": "hi", "run_id": "r1"}, - ) - await proxy_client.post("/session/s1/close") - h = await proxy_client.get("/session/s1/history") - assert h.json()["history"] == [] - - -class TestRouterIntegration: - @pytest.mark.asyncio - async def test_register_and_route(self): - router_app = create_router_app( - RouterConfig(admin_api_key=DEFAULT_ADMIN_API_KEY) - ) - transport = httpx.ASGITransport(app=router_app) - - async with httpx.AsyncClient( - transport=transport, base_url="http://router" - ) as client: - await client.post( - "/register", - json={"addr": "http://proxy1:9100"}, - headers=_AUTH, - ) - resp = await client.post( - "/route", json={"session_key": "s1"}, headers=_AUTH - ) - assert resp.json()["data_proxy_addr"] == "http://proxy1:9100" - - resp2 = await client.post( - "/route", json={"session_key": "s1"}, headers=_AUTH - ) - assert resp2.json()["data_proxy_addr"] == "http://proxy1:9100" - - -class TestToolCallFlow: - @pytest.mark.asyncio - async def test_tool_events_through_proxy(self): - worker_app = _make_worker_app(_ToolAgent) - worker_transport = httpx.ASGITransport(app=worker_app) - proxy_app = create_data_proxy_app(DataProxyConfig(worker_addr="http://worker")) - - original_post = httpx.AsyncClient.post - - async def patched_post(self, url, **kwargs): - if "worker" in url: - async with httpx.AsyncClient( - transport=worker_transport, base_url="http://worker" - ) as wc: - path = url.split("http://worker")[-1] - return await wc.post(path, **kwargs) - return await original_post(self, url, **kwargs) - - proxy_transport = httpx.ASGITransport(app=proxy_app) - - with patch.object(httpx.AsyncClient, "post", patched_post): - async with httpx.AsyncClient( - transport=proxy_transport, base_url="http://proxy" - ) as proxy_client: - resp = await proxy_client.post( - "/session/s1/turn", - json={"message": "lookup 123", "run_id": "r1"}, - ) - data = resp.json() - assert data["summary"] == "Lookup complete" - events = data["events"] - types = {e["type"] for e in events} - assert "tool_call" in types - assert "tool_result" in types - - # History should include tool call records - h = await proxy_client.get("/session/s1/history") - history = h.json()["history"] - tool_msgs = [m for m in history if m.get("role") == "tool"] - assert len(tool_msgs) > 0 - assert "tool_call_id" in tool_msgs[0] - - -class TestGatewayHealth: - @pytest.mark.asyncio - async def test_health(self): - app = create_gateway_app(GatewayConfig(router_addr="http://fake-router")) - transport = httpx.ASGITransport(app=app) - async with httpx.AsyncClient( - transport=transport, base_url="http://gw" - ) as client: - resp = await client.get("/health") - assert resp.status_code == 200 - assert resp.json()["status"] == "ok" - - -class TestBridgeExtractMessage: - def test_text_message(self): - items = [ - { - "type": "message", - "content": [{"type": "input_text", "text": "Hello"}], - } - ] - assert OpenResponsesBridge._extract_message(items, "") == "Hello" - - def test_string_content(self): - items = [{"type": "message", "content": "Simple"}] - assert OpenResponsesBridge._extract_message(items, "") == "Simple" - - def test_instructions_prepended(self): - items = [{"type": "message", "content": "Hi"}] - result = OpenResponsesBridge._extract_message(items, "Be helpful") - assert result.startswith("Be helpful") - assert "Hi" in result - - def test_function_call_output(self): - items = [{"type": "function_call_output", "output": "42"}] - result = OpenResponsesBridge._extract_message(items, "") - assert "[tool result] 42" in result - - -class TestBridgeDeriveSessionKey: - def test_with_user(self): - key = OpenResponsesBridge._derive_session_key("user1", "model1") - assert key == "agent:model1:user1" - - def test_without_user_is_unique(self): - k1 = OpenResponsesBridge._derive_session_key("", "m") - k2 = OpenResponsesBridge._derive_session_key("", "m") - assert k1 != k2 - assert k1.startswith("agent:m:") - - def test_default_model(self): - key = OpenResponsesBridge._derive_session_key("u1", "") - assert key == "agent:default:u1" diff --git a/tests/experimental/openai/test_client_with_tool.py b/tests/experimental/openai/test_client_with_tool.py index f53dfdddb5..50c0465cdb 100644 --- a/tests/experimental/openai/test_client_with_tool.py +++ b/tests/experimental/openai/test_client_with_tool.py @@ -232,16 +232,23 @@ async def test_streaming_with_tool_calling(openai_client): and len(chunk.choices[0].delta.tool_calls) > 0 ] - # Verify tool calls are present and properly structured + # Verify tool calls are present and split into name/id and arguments chunks. if tool_call_chunks: - for chunk in tool_call_chunks: - for tool_call in chunk.choices[0].delta.tool_calls: - # Verify tool_call has required fields (simplified access after our fix) - assert tool_call.id is not None - assert tool_call.type == "function" - assert tool_call.function is not None - assert tool_call.function.name is not None - assert tool_call.function.arguments is not None + assert len(tool_call_chunks) % 2 == 0 + for first_chunk, second_chunk in zip( + tool_call_chunks[::2], tool_call_chunks[1::2], strict=True + ): + first_tool_call = first_chunk.choices[0].delta.tool_calls[0] + second_tool_call = second_chunk.choices[0].delta.tool_calls[0] + assert first_tool_call.id is not None + assert first_tool_call.type == "function" + assert first_tool_call.function is not None + assert first_tool_call.function.name is not None + assert first_tool_call.function.arguments == "" + assert second_tool_call.index == first_tool_call.index + assert second_tool_call.function is not None + assert second_tool_call.function.arguments is not None + assert len(second_tool_call.function.arguments) > 0 # Find the final chunk with finish_reason final_chunks = [ diff --git a/tests/experimental/openai/test_parse_tool_call_arguments.py b/tests/experimental/openai/test_parse_tool_call_arguments.py new file mode 100644 index 0000000000..afb702ebcc --- /dev/null +++ b/tests/experimental/openai/test_parse_tool_call_arguments.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 + +from areal.experimental.openai.client import _parse_tool_call_arguments + + +def test_string_arguments_are_parsed_to_dict(): + messages = [ + { + "role": "assistant", + "tool_calls": [ + {"function": {"name": "Bash", "arguments": '{"command": "ls"}'}} + ], + } + ] + + out = _parse_tool_call_arguments(messages) + + assert out[0]["tool_calls"][0]["function"]["arguments"] == {"command": "ls"} + + +def test_non_string_arguments_are_left_unchanged(): + messages = [ + { + "role": "assistant", + "tool_calls": [ + {"function": {"name": "Bash", "arguments": {"command": "ls"}}} + ], + } + ] + + out = _parse_tool_call_arguments(messages) + + assert out[0] is messages[0] + assert out[0]["tool_calls"][0]["function"]["arguments"] == {"command": "ls"} + + +def test_invalid_json_arguments_are_left_unchanged(): + messages = [ + { + "role": "assistant", + "tool_calls": [{"function": {"name": "Bash", "arguments": "not json"}}], + } + ] + + out = _parse_tool_call_arguments(messages) + + assert out[0]["tool_calls"][0]["function"]["arguments"] == "not json" + + +def test_messages_without_tool_calls_pass_through(): + messages = [{"role": "user", "content": "hi"}] + + out = _parse_tool_call_arguments(messages) + + assert out[0] is messages[0] diff --git a/tests/experimental/openai/test_resolve_max_total_tokens.py b/tests/experimental/openai/test_resolve_max_total_tokens.py new file mode 100644 index 0000000000..826e6f2062 --- /dev/null +++ b/tests/experimental/openai/test_resolve_max_total_tokens.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 + +from areal.experimental.openai.client import ( + _DEFAULT_MAX_TOTAL_TOKENS, + _resolve_max_total_tokens, +) + + +def test_engine_max_tokens_caps_total_length(): + out = _resolve_max_total_tokens( + prompt_len=1000, + max_new_tokens=100_000, + engine_max_tokens=4096, + ) + + assert out == 4096 + + +def test_fallback_ceiling_applies_when_engine_max_tokens_is_none(): + out = _resolve_max_total_tokens( + prompt_len=1000, + max_new_tokens=100_000, + engine_max_tokens=None, + ) + + assert out == _DEFAULT_MAX_TOTAL_TOKENS + + +def test_returns_prompt_plus_generation_when_under_cap(): + out = _resolve_max_total_tokens( + prompt_len=100, + max_new_tokens=200, + engine_max_tokens=4096, + ) + + assert out == 300 diff --git a/tests/experimental/openai/test_tool_call_parser.py b/tests/experimental/openai/test_tool_call_parser.py index 4d3da68501..c08c4aebe3 100644 --- a/tests/experimental/openai/test_tool_call_parser.py +++ b/tests/experimental/openai/test_tool_call_parser.py @@ -1,3 +1,4 @@ +import json from types import SimpleNamespace import pytest @@ -69,6 +70,52 @@ '\n{"name": "search", "arguments": {"query": "director of The Preacher\'s Wife country"}}\n<|im_end|>' ) +QWEN3_CODER_TOOLS: list[ChatCompletionToolParam] = [ + { + "type": "function", + "function": { + "name": "Bash", + "description": "Run a shell command.", + "parameters": { + "type": "object", + "properties": { + "command": {"type": "string"}, + "timeout": {"type": "integer"}, + }, + "required": ["command"], + }, + }, + }, + { + "type": "function", + "function": { + "name": "Write", + "description": "Write a file.", + "parameters": { + "type": "object", + "properties": { + "file_path": {"type": "string"}, + "content": {"type": "string"}, + }, + "required": ["file_path", "content"], + }, + }, + }, +] + +QWEN3_CODER_TEXT = """I'll inspect the repository. + + + + +ls -la + + +120 + + +<|im_end|>""" + def _assert_tool_calls(tool_calls, new_text: str, new_finish_reason: str) -> None: assert new_finish_reason == "tool_calls" @@ -101,6 +148,67 @@ def _run_process_tool_calls(text: str): ) +def test_process_tool_calls_qwen3_coder_xml_parameters_without_sglang(): + tool_calls, new_text, new_finish_reason = parser_module.process_tool_calls( + text=QWEN3_CODER_TEXT, + tools=QWEN3_CODER_TOOLS, + tool_call_parser="qwen3_coder", + reasoning_parser="qwen3", + finish_reason="stop", + use_responses=False, + tokenizer=object(), + ) + + assert new_finish_reason == "tool_calls" + assert tool_calls is not None + assert len(tool_calls) == 1 + assert tool_calls[0].function.name == "Bash" + assert json.loads(tool_calls[0].function.arguments) == { + "command": "ls -la", + "timeout": 120, + } + assert "" not in new_text + assert "<|im_end|>" not in new_text + assert "I'll inspect the repository." in new_text + + +def test_qwen3_coder_xml_empty_args_returns_no_parse(): + text = "\n\n\n" + + tool_calls, new_text, finish_reason = ( + parser_module._process_tool_calls_qwen3_coder_xml( + text=text, + tools=QWEN3_CODER_TOOLS, + finish_reason="stop", + ) + ) + + assert tool_calls is None + assert new_text == text + assert finish_reason == "stop" + + +def test_qwen3_coder_xml_literal_closing_tag_is_not_silently_truncated(): + text = ( + "\n\n" + "\na.xml\n\n" + "\nsome XML text: more content\n\n" + "\n" + ) + + tool_calls, new_text, finish_reason = ( + parser_module._process_tool_calls_qwen3_coder_xml( + text=text, + tools=QWEN3_CODER_TOOLS, + finish_reason="stop", + ) + ) + + assert tool_calls is None + assert new_text == text + assert finish_reason == "stop" + + @pytest.mark.sglang def test_process_tool_calls_qwen25_chat_completions_sglang(): pytest.importorskip( diff --git a/tests/infra/test_launcher_utils.py b/tests/infra/test_launcher_utils.py new file mode 100644 index 0000000000..1f2a4601f6 --- /dev/null +++ b/tests/infra/test_launcher_utils.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +import subprocess +from types import SimpleNamespace +from unittest.mock import patch + +from areal.infra.utils.launcher import run_post_exit_hook + + +def _config(post_exit_hook: str): + return SimpleNamespace( + post_exit_hook=post_exit_hook, + cluster=SimpleNamespace(fileroot="/tmp/areal-test"), + experiment_name="exp", + trial_name="trial", + ) + + +def test_run_post_exit_hook_empty_hook_skips_subprocess(): + with patch("areal.infra.utils.launcher.subprocess.run") as mock_run: + run_post_exit_hook(_config("")) + + mock_run.assert_not_called() + + +def test_run_post_exit_hook_command_injects_log_dir(): + with patch("areal.infra.utils.launcher.getpass.getuser", return_value="tester"): + with patch("areal.infra.utils.launcher.subprocess.run") as mock_run: + mock_run.return_value = subprocess.CompletedProcess( + args="cleanup", + returncode=0, + stdout="ok", + stderr="", + ) + + run_post_exit_hook(_config(" cleanup ")) + + mock_run.assert_called_once() + _, kwargs = mock_run.call_args + assert mock_run.call_args.args[0] == "cleanup" + assert kwargs["shell"] is True + assert kwargs["timeout"] == 600 + assert kwargs["capture_output"] is True + assert kwargs["text"] is True + assert kwargs["env"]["LOG_DIR"] == "/tmp/areal-test/logs/tester/exp/trial" + + +def test_run_post_exit_hook_timeout_does_not_raise(): + with patch("areal.infra.utils.launcher.subprocess.run") as mock_run: + mock_run.side_effect = subprocess.TimeoutExpired( + cmd="cleanup", + timeout=600, + ) + + run_post_exit_hook(_config("cleanup")) + + mock_run.assert_called_once() diff --git a/tests/test_adv_norm_config.py b/tests/test_adv_norm_config.py index fb1553fee4..6e9adf8b14 100644 --- a/tests/test_adv_norm_config.py +++ b/tests/test_adv_norm_config.py @@ -874,11 +874,15 @@ def test_group_size_edge_cases(): adv_norm = Normalization(config) normalized = adv_norm(advantages, loss_mask) - # With group size 1 and leave-one-out, each element should remain approximately unchanged - # because mean=0 (no other elements) and std=1 (for stability), so result ≈ (x-0)/1 = x + + # With singleton leave-one-out groups, there is no peer baseline. + # We use the sample itself as the baseline, so each singleton group + # normalizes to zero instead of passing through the raw advantage. assert torch.allclose( - normalized, advantages, atol=1e-3 - ) # Should remain approximately unchanged + normalized, + torch.zeros_like(advantages), + atol=1e-6, + ) def test_precision_and_numerical_stability(): diff --git a/tests/test_cispo_loss.py b/tests/test_cispo_loss.py new file mode 100644 index 0000000000..44745a5a31 --- /dev/null +++ b/tests/test_cispo_loss.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CISPO surrogate (MiniMax-M1 Eq. 4-5) invariants. + +Two defining properties, each across a PPO-like band (0.2 / 0.28) and the wide +MiniMax-M1 band (1.0 / 4.0): + +1. closed-form surrogate value + which tokens the clip flags; +2. gradient routing -- ``logprobs.grad == -sg(clip(ratio)) * A / N`` exactly, + with zero gradient through the importance-ratio path (the test fails if the + stop-gradient detach is dropped). +""" + +import pytest +import torch + +from areal.api.cli_args import PPOActorConfig, RejectionSamplingConfig +from areal.utils.functional import cispo_loss_fn + +# (eps_clip, eps_clip_higher): a PPO-like band that clips most tokens, and the +# wide MiniMax-M1 band [0, 5] that clips none of the fixture ratios. +BANDS = [(0.2, 0.28), (1.0, 4.0)] + + +def _inputs(): + # log_ratio spans below / inside / above a tight band. The last token is + # masked out, so it must contribute neither loss nor gradient. + log_ratio = torch.tensor([-2.0, -0.5, 0.0, 0.3, 1.5, 0.1]) + logprobs = torch.tensor([-1.0, -2.0, -0.3, -1.2, -0.7, -2.5], requires_grad=True) + proximal = logprobs.detach() - log_ratio # logprobs - proximal == log_ratio + advantages = torch.tensor([1.0, -2.0, 0.5, 3.0, -1.0, 2.0]) + loss_mask = torch.tensor([1, 1, 1, 1, 1, 0], dtype=torch.bool) + return logprobs, proximal, advantages, loss_mask, log_ratio + + +@pytest.mark.parametrize("eps_clip,eps_clip_higher", BANDS) +def test_cispo_closed_form_value_and_clip_mask(eps_clip, eps_clip_higher): + logprobs, proximal, advantages, loss_mask, log_ratio = _inputs() + + loss, stat = cispo_loss_fn( + logprobs=logprobs, + proximal_logprobs=proximal, + advantages=advantages, + eps_clip=eps_clip, + eps_clip_higher=eps_clip_higher, + loss_mask=loss_mask, + ) + + ratio = torch.exp(log_ratio) + ratio_clipped = ratio.clamp(1.0 - eps_clip, 1.0 + eps_clip_higher) + per_token = -ratio_clipped * advantages * logprobs.detach() + expected = ( + torch.where(loss_mask, per_token, torch.zeros_like(per_token)).sum() + / loss_mask.count_nonzero() + ) + + torch.testing.assert_close(loss, expected) + # importance_weight logs the *unclipped* ratio for diagnostics. + torch.testing.assert_close(stat["importance_weight"], ratio) + # clip_mask flags band-exit on either side, intersected with the loss mask. + expected_clip = (ratio != ratio_clipped) & loss_mask + assert torch.equal(stat["clip_mask"], expected_clip) + assert not stat["dual_clip_mask"].any() + + +@pytest.mark.parametrize("eps_clip,eps_clip_higher", BANDS) +def test_cispo_gradient_routes_through_logprobs_only(eps_clip, eps_clip_higher): + logprobs, proximal, advantages, loss_mask, log_ratio = _inputs() + # Make the proximal logp a grad-tracking leaf: a correct stop-gradient must + # leave it with no gradient. + proximal = proximal.clone().requires_grad_(True) + + loss, _ = cispo_loss_fn( + logprobs=logprobs, + proximal_logprobs=proximal, + advantages=advantages, + eps_clip=eps_clip, + eps_clip_higher=eps_clip_higher, + loss_mask=loss_mask, + ) + loss.backward() + + ratio = torch.exp(log_ratio) + ratio_clipped = ratio.clamp(1.0 - eps_clip, 1.0 + eps_clip_higher) + n = loss_mask.count_nonzero() + # Gradient flows ONLY through the explicit `logprobs` factor; the clipped + # ratio is a stop-gradient constant. If the detach were dropped, ratio would + # depend on logprobs and add a term, breaking this exact equality. + expected_grad = torch.where( + loss_mask, -ratio_clipped * advantages / n, torch.zeros_like(ratio) + ) + torch.testing.assert_close(logprobs.grad, expected_grad) + # No gradient leaks to the importance-ratio (proximal) path. + assert proximal.grad is None + + +def test_cispo_rejects_nonpositive_eps_clip_higher(): + logprobs, proximal, advantages, loss_mask, _ = _inputs() + for bad in (None, 0.0, -1.0): + with pytest.raises(ValueError, match="eps_clip_higher"): + cispo_loss_fn( + logprobs=logprobs, + proximal_logprobs=proximal, + advantages=advantages, + eps_clip=1.0, + eps_clip_higher=bad, + loss_mask=loss_mask, + ) + + +def test_cispo_decoupled_applies_behave_imp_weight(): + # Decoupled loss: surrogate rescaled by detached pi_proximal/pi_behave; grad + # still routes only through logprobs. + eps_clip, eps_clip_higher = 1.0, 4.0 + logprobs, proximal, advantages, loss_mask, log_ratio = _inputs() + # upper is wide so clamp is a no-op and behave_imp_weight == exp(behave_log_ratio). + behave_log_ratio = torch.tensor([0.2, -0.4, 0.0, 0.1, -0.3, 0.5]) + old_logprobs = proximal - behave_log_ratio # proximal - old == behave_log_ratio + rs = RejectionSamplingConfig( + level="token", action="clamp", metric="ratio", upper=100.0 + ) + + proximal_leaf = proximal.clone().requires_grad_(True) + old_leaf = old_logprobs.clone().requires_grad_(True) + + loss, stat = cispo_loss_fn( + logprobs=logprobs, + proximal_logprobs=proximal_leaf, + advantages=advantages, + eps_clip=eps_clip, + eps_clip_higher=eps_clip_higher, + loss_mask=loss_mask, + old_logprobs=old_leaf, + rejection_sampling=rs, + ) + + ratio = torch.exp(log_ratio) + ratio_clipped = ratio.clamp(1.0 - eps_clip, 1.0 + eps_clip_higher) + behave_w = torch.exp(behave_log_ratio) + n = loss_mask.count_nonzero() + per_token = -ratio_clipped * advantages * logprobs.detach() * behave_w + expected = torch.where(loss_mask, per_token, torch.zeros_like(per_token)).sum() / n + torch.testing.assert_close(loss, expected) + # behave_imp_weight is zeroed on masked tokens, so compare valid positions. + torch.testing.assert_close( + stat["behave_imp_weight"][loss_mask], behave_w[loss_mask] + ) + assert "behave_mask" in stat and "behave_approx_kl" in stat + + loss.backward() + expected_grad = torch.where( + loss_mask, -ratio_clipped * advantages * behave_w / n, torch.zeros_like(ratio) + ) + torch.testing.assert_close(logprobs.grad, expected_grad) + assert proximal_leaf.grad is None + assert old_leaf.grad is None + + +def test_cispo_config_validation(): + # Requires a positive upper clip. + with pytest.raises(ValueError, match="eps_clip_higher"): + PPOActorConfig(use_cispo_loss=True, eps_clip_higher=None) + # Mutually exclusive with SAPO. + with pytest.raises(ValueError, match="mutually exclusive"): + PPOActorConfig(use_cispo_loss=True, use_sapo_loss=True, eps_clip_higher=4.0) + # Token level only. + with pytest.raises(ValueError, match="importance_sampling_level"): + PPOActorConfig( + use_cispo_loss=True, + eps_clip_higher=4.0, + importance_sampling_level="sequence", + ) + # Valid configuration does not raise. + PPOActorConfig(use_cispo_loss=True, eps_clip=1.0, eps_clip_higher=4.0) diff --git a/tests/test_clevr_reward.py b/tests/test_clevr_reward.py new file mode 100644 index 0000000000..9d9c0b4d8d --- /dev/null +++ b/tests/test_clevr_reward.py @@ -0,0 +1,15 @@ +from areal.reward import clevr_count_70k_reward_fn + + +def test_clevr_non_string_answer_scored_correctly(): + """A non-string answer (e.g. an int) must be coerced and scored, not crash. + The sibling reward fns str()-coerce both inputs; clevr did not, so a matching + completion scored as a dropped trajectory instead of 1.0.""" + reward = clevr_count_70k_reward_fn( + prompt="", + completions="[3]", + prompt_ids=[], + completion_ids=[], + answer=3, + ) + assert reward == 1.0 diff --git a/tests/test_examples.py b/tests/test_examples.py index 2acb62f18a..7b8cf41805 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -306,6 +306,7 @@ def test_vlm_grpo(tmp_path_factory, rollout_backend, actor_backend): "gconfig.n_samples=2", "gconfig.max_new_tokens=256", "actor.mb_spec.max_tokens_per_mb=1024", + "+actor.optimizer_dtype=bfloat16", "train_dataset.batch_size=2", "valid_dataset.batch_size=2", f"train_dataset.path={dataset_path}", @@ -814,7 +815,7 @@ def test_tau2(tmp_path_factory): [ "python3", "-m", - "areal.experimental.inference_service.sglang.launch_server", + "areal.v2.inference_service.sglang.launch_server", "--model-path", model_path, "--host", diff --git a/tests/test_megatron_async_save.py b/tests/test_megatron_async_save.py index 93044464cf..b7041c0e83 100644 --- a/tests/test_megatron_async_save.py +++ b/tests/test_megatron_async_save.py @@ -260,3 +260,42 @@ def test_sync_save_emits_no_async_metrics(patched_checkpointer, tmp_path): manager.save_checkpoint(str(tmp_path / "step0")) stats_scalar.assert_not_called() + + +def test_generate_state_dict_requests_dp_reshardable_sharding(patched_checkpointer): + _, manager, _ = patched_checkpointer + + with patch("torch.distributed.barrier"): + state_dict = manager.generate_state_dict( + with_model=False, with_optimizer=True, with_rng=False + ) + + kwargs = manager.optimizer.sharded_state_dict.call_args.kwargs + assert kwargs["metadata"] == {"distrib_optim_sharding_type": "dp_reshardable"} + assert kwargs["is_loading"] is False + assert "optimizer" in state_dict + + +def test_load_checkpoint_builds_optimizer_template_with_is_loading( + patched_checkpointer, tmp_path +): + mod, manager, _ = patched_checkpointer + + with ( + patch("os.path.exists", return_value=True), + patch("torch.distributed.barrier"), + patch.object( + mod, "load_dist_checkpointing", return_value={"optimizer": {"step": 1}} + ), + ): + manager.load_checkpoint( + str(tmp_path / "step0"), + with_model=False, + with_optimizer=True, + with_rng=False, + ) + + kwargs = manager.optimizer.sharded_state_dict.call_args.kwargs + assert kwargs["is_loading"] is True + assert kwargs["metadata"] == {"distrib_optim_sharding_type": "dp_reshardable"} + manager.optimizer.load_state_dict.assert_called_once_with({"step": 1}) diff --git a/tests/test_network_ipv6_utils.py b/tests/test_network_ipv6_utils.py index 273105d957..235ceba311 100644 --- a/tests/test_network_ipv6_utils.py +++ b/tests/test_network_ipv6_utils.py @@ -86,3 +86,13 @@ def test_get_loopback_ip_returns_bindable_address(): network = _load_network_module() ip = network.get_loopback_ip() assert ip in {"127.0.0.1", "::1"} + + +def test_find_free_ports_ignores_out_of_range_excludes(): + # Excluded ports outside [min_port, max_port] must not deflate the + # availability check; all in-range ports here are free. + network = _load_network_module() + ports = network.find_free_ports( + count=6, port_range=(10000, 10005), exclude_ports={50000} + ) + assert len(ports) == 6 diff --git a/tests/test_packed_vs_padded_consistency.py b/tests/test_packed_vs_padded_consistency.py index 7e1cee73a5..7dc08ab4ee 100644 --- a/tests/test_packed_vs_padded_consistency.py +++ b/tests/test_packed_vs_padded_consistency.py @@ -69,6 +69,7 @@ def test_llm_consistency(model_path, mock_padded_llm_data): backend="fsdp:d1", path=model_path, dtype="bfloat16", + optimizer_dtype="bfloat16", attn_impl="flash_attention_2", gradient_checkpointing=False, disable_dropout=True, diff --git a/tests/test_recover.py b/tests/test_recover.py index 57e5634cd5..653e63a004 100644 --- a/tests/test_recover.py +++ b/tests/test_recover.py @@ -7,10 +7,10 @@ from areal.api.cli_args import RecoverConfig from areal.api.io_struct import FinetuneSpec, StepInfo -from areal.experimental.training_service.controller.controller import ( +from areal.utils.recover import RecoverHandler, check_if_auto_recover, check_if_recover +from areal.v2.training_service.controller.controller import ( GatewayTrainController, ) -from areal.utils.recover import RecoverHandler, check_if_auto_recover, check_if_recover class TestRecoverConfig: diff --git a/tests/test_reuse_train_logp.py b/tests/test_reuse_train_logp.py new file mode 100644 index 0000000000..4b82747711 --- /dev/null +++ b/tests/test_reuse_train_logp.py @@ -0,0 +1,37 @@ +"""Tests for the ``reuse_train_logp`` proximal-logp method. + +``reuse_train_logp`` reuses the training forward-pass logprobs as the proximal +logp, skipping the extra decoupled-PPO forward. It requires +``ppo_n_minibatches == 1`` so the training forward still reflects the policy +that produced the rollout (with multiple minibatches the weights change between +steps, so the reused logprobs would no longer be the proximal policy). +""" + +from areal.api.cli_args import PPOActorConfig +from areal.utils.constants import ( + PROX_LOGP_METHOD_REUSE_TRAIN_LOGP, + PROX_LOGP_METHODS_ALL, + ProxLogpMethod, +) + + +def test_enum_includes_reuse_train_logp(): + assert ProxLogpMethod.REUSE_TRAIN_LOGP.value == "reuse_train_logp" + assert ProxLogpMethod("reuse_train_logp") is ProxLogpMethod.REUSE_TRAIN_LOGP + assert "reuse_train_logp" in PROX_LOGP_METHODS_ALL + assert PROX_LOGP_METHOD_REUSE_TRAIN_LOGP == "reuse_train_logp" + + +def test_reuse_train_logp_skips_forward_pass(): + # It must skip the extra proximal forward, like loglinear does. + assert ProxLogpMethod.REUSE_TRAIN_LOGP.skips_forward_pass() + + +def test_reuse_train_logp_allows_single_minibatch(): + config = PPOActorConfig( + backend="fsdp:d1", + prox_logp_method="reuse_train_logp", + ppo_n_minibatches=1, + ) + assert config.prox_logp_method == "reuse_train_logp" + assert config.ppo_n_minibatches == 1 diff --git a/tests/test_reward_norm_variable_group.py b/tests/test_reward_norm_variable_group.py new file mode 100644 index 0000000000..c6ce65232f --- /dev/null +++ b/tests/test_reward_norm_variable_group.py @@ -0,0 +1,185 @@ +"""Tests for variable-size group support in reward/advantage normalization. + +These cover the ``group_sizes`` argument added to ``Normalization`` so that +group-level normalization slices each trajectory group by its actual sample +count instead of a fixed ``group_size``. Variable-size groups arise when some +rollout samples fail or are filtered, leaving groups smaller than ``n_samples``. +A fixed-size slice would otherwise straddle two groups or leave a tail. +""" + +import pytest +import torch + +from areal.api.cli_args import NormConfig +from areal.utils.data import Normalization, batched_call + + +def _group_norm_config() -> NormConfig: + return NormConfig( + mean_level="group", + std_level="group", + group_size=2, + mean_leave1out=False, + std_unbiased=False, + eps=0.0, + ) + + +def _ref_group_norm(x: torch.Tensor, boundaries: list[int]) -> torch.Tensor: + """Reference: independently mean/std normalize each variable-size group.""" + out = torch.empty_like(x, dtype=torch.float64) + offset = 0 + for size in boundaries: + chunk = x[offset : offset + size].to(torch.float64) + mean = chunk.mean() + std = chunk.std(unbiased=False) + out[offset : offset + size] = (chunk - mean) / std if std > 0 else 0.0 + offset += size + return out + + +def test_variable_group_sizes_normalize_each_group_independently(): + """Unequal group sizes are sliced exactly, not by a fixed group_size.""" + norm = Normalization(_group_norm_config()) + # Three groups of sizes 3, 2, 3 (total 8); fixed group_size=2 would be wrong. + x = torch.tensor( + [0.0, 3.0, 6.0, 10.0, 20.0, 1.0, 2.0, 3.0], + dtype=torch.float32, + ) + boundaries = [3, 2, 3] + + out = norm(x, group_sizes=boundaries) + expected = _ref_group_norm(x, boundaries) + + torch.testing.assert_close(out.double(), expected, rtol=1e-5, atol=1e-5) + # Each group is zero-mean after normalization. + offset = 0 + for size in boundaries: + torch.testing.assert_close( + out[offset : offset + size].double().mean(), + torch.tensor(0.0, dtype=torch.float64), + rtol=1e-5, + atol=1e-5, + ) + offset += size + + +def test_group_sizes_none_matches_fixed_group_size(): + """Without group_sizes, behavior is unchanged (backward compatible).""" + norm = Normalization(_group_norm_config()) # group_size=2 + x = torch.tensor([0.0, 2.0, 10.0, 14.0], dtype=torch.float32) + + out_default = norm(x) + out_explicit = norm(x, group_sizes=[2, 2]) + + torch.testing.assert_close(out_default, out_explicit, rtol=1e-6, atol=1e-6) + + +def test_variable_group_with_zero_variance_group_is_finite(): + """A zero-variance group stays finite thanks to eps in the denominator. + + When a group's rewards are all equal its std is 0; dividing by ``std + eps`` + (eps > 0) keeps the result finite (and zero, since the numerator is also 0) + instead of producing NaN/inf. This is the guard that prevents a degenerate + group from blowing up advantages. + """ + cfg = NormConfig( + mean_level="group", + std_level="group", + group_size=2, + mean_leave1out=False, + std_unbiased=False, + eps=1e-5, + ) + norm = Normalization(cfg) + # Group 0 (size 3) has variance; group 1 (size 2) is all-equal -> std 0. + x = torch.tensor([0.0, 3.0, 6.0, 5.0, 5.0], dtype=torch.float32) + boundaries = [3, 2] + + out = norm(x, group_sizes=boundaries) + + assert torch.isfinite(out).all(), "normalized output must be finite" + # The degenerate group collapses to zero (x_centered=0 / (0+eps)). + torch.testing.assert_close( + out[3:].double(), + torch.zeros(2, dtype=torch.float64), + rtol=0, + atol=0, + ) + + +def test_leave_one_out_singleton_group_outputs_zero(): + """A size-1 leave-one-out group has no peer baseline, so it is zeroed.""" + cfg = NormConfig( + mean_level="group", + std_level="group", + group_size=1, + mean_leave1out=True, + std_unbiased=True, + eps=1e-5, + ) + norm = Normalization(cfg) + x = torch.tensor([7.0, 1.0, 5.0], dtype=torch.float32) + # One singleton group then a size-2 group. + out = norm(x, group_sizes=[1, 2]) + + assert torch.isfinite(out).all() + assert out[0].item() == 0.0 + + +def test_group_sizes_sum_mismatch_raises(): + """Boundaries whose sizes do not sum to the batch size are rejected.""" + norm = Normalization(_group_norm_config()) + x = torch.tensor([0.0, 1.0, 2.0, 3.0], dtype=torch.float32) # bs=4 + with pytest.raises(ValueError, match="must equal"): + norm(x, group_sizes=[3, 2]) # sums to 5 != 4 + + +def test_group_sizes_non_positive_raises(): + """Boundaries containing a non-positive size are rejected.""" + norm = Normalization(_group_norm_config()) + x = torch.tensor([0.0, 1.0, 2.0, 3.0], dtype=torch.float32) # bs=4 + with pytest.raises(ValueError, match="positive"): + norm(x, group_sizes=[4, 0]) + + +def test_adv_norm_style_2d_variable_groups_normalize_per_group(): + """Token-level (2D) advantages normalize per variable-size group on dim 0. + + Mirrors how PPOActor.adv_norm is applied to the GAE advantage tensor. + """ + norm = _group_norm_config() # group mean+std, eps=0 + adv_norm = Normalization(norm) + # bs=5 (groups [3, 2]), seqlen=2. + x = torch.tensor( + [[0.0, 6.0], [3.0, 6.0], [6.0, 6.0], [10.0, 0.0], [20.0, 0.0]], + dtype=torch.float32, + ) + loss_mask = torch.ones_like(x) + out = adv_norm(x, loss_mask, group_sizes=[3, 2]) + + assert torch.isfinite(out).all() + assert out.shape == x.shape + + +def test_batched_call_pass_meta_provides_traj_group_sizes(): + data = [ + { + "attention_mask": torch.ones(2, 3, dtype=torch.bool), + "values": torch.tensor([[1.0], [2.0]]), + }, + { + "attention_mask": torch.ones(1, 3, dtype=torch.bool), + "values": torch.tensor([[3.0]]), + }, + ] + + def fn(batch, meta): + assert meta.traj_group_sizes == [2, 1] + return {"values": batch["values"] + 1} + + out = batched_call(fn, data, pass_meta=True) + + assert len(out) == 2 + torch.testing.assert_close(out[0]["values"], torch.tensor([[2.0], [3.0]])) + torch.testing.assert_close(out[1]["values"], torch.tensor([[4.0]])) diff --git a/tests/test_sglang_pp_unit.py b/tests/test_sglang_pp_unit.py index a9afe74263..203023c124 100644 --- a/tests/test_sglang_pp_unit.py +++ b/tests/test_sglang_pp_unit.py @@ -313,14 +313,14 @@ class TestPPBridgeModule: """Test the pp_bridge module can be imported and has expected symbols.""" def test_pp_bridge_class_exists(self): - from areal.experimental.inference_service.sglang.pp_bridge import ( + from areal.v2.inference_service.sglang.pp_bridge import ( PPSchedulerBridge, ) assert PPSchedulerBridge is not None def test_extract_pp_rank_from_group_name(self): - from areal.experimental.inference_service.sglang.pp_bridge import ( + from areal.v2.inference_service.sglang.pp_bridge import ( _extract_pp_rank_from_group_name, ) @@ -331,7 +331,7 @@ def test_extract_pp_rank_from_group_name(self): assert _extract_pp_rank_from_group_name("update_weight_group") is None def test_pp_bridge_bind_is_callable(self): - from areal.experimental.inference_service.sglang.pp_bridge import ( + from areal.v2.inference_service.sglang.pp_bridge import ( PPSchedulerBridge, ) @@ -339,7 +339,7 @@ def test_pp_bridge_bind_is_callable(self): def test_pp_bridge_noop_when_pp1(self): """PPSchedulerBridge.bind() should be a no-op when pp_size <= 1.""" - from areal.experimental.inference_service.sglang.pp_bridge import ( + from areal.v2.inference_service.sglang.pp_bridge import ( PPSchedulerBridge, ) diff --git a/tests/test_split_trajectory_for_dump.py b/tests/test_split_trajectory_for_dump.py new file mode 100644 index 0000000000..8b1907edb2 --- /dev/null +++ b/tests/test_split_trajectory_for_dump.py @@ -0,0 +1,119 @@ +"""Tests for WorkflowExecutor._split_trajectory_for_dump and _compute_output_versions.""" + +import pytest + +from areal.infra.workflow_executor import WorkflowExecutor + + +class _FakeTokenizer: + """Minimal tokenizer that returns a deterministic string for testing.""" + + def decode(self, ids: list[int], **kwargs) -> str: + return f"[{len(ids)} tokens]" + + +class TestSplitTrajectoryForDump: + @pytest.fixture + def tokenizer(self): + return _FakeTokenizer() + + def test_single_turn_no_segments(self, tokenizer): + ids = [1, 2, 3, 4, 5, 6] + mask = [0, 0, 0, 1, 1, 1] + result = WorkflowExecutor._split_trajectory_for_dump(ids, mask, tokenizer) + assert result["prompt_end"] == 3 + assert result["prompt_text"] == "[3 tokens]" + assert result["completion_text"] == "[3 tokens]" + assert result["segments"] is None + + def test_multi_turn_has_segments(self, tokenizer): + ids = list(range(8)) + mask = [0, 0, 1, 1, 0, 0, 1, 1] + result = WorkflowExecutor._split_trajectory_for_dump(ids, mask, tokenizer) + assert result["prompt_end"] == 2 + assert result["segments"] is not None + assert len(result["segments"]) == 4 + roles = [s["role"] for s in result["segments"]] + assert roles == ["prompt", "gen", "context", "gen"] + lengths = [s["len"] for s in result["segments"]] + assert lengths == [2, 2, 2, 2] + + def test_all_zeros_prompt_only(self, tokenizer): + ids = [1, 2, 3] + mask = [0, 0, 0] + result = WorkflowExecutor._split_trajectory_for_dump(ids, mask, tokenizer) + assert result["prompt_end"] == 3 + assert result["prompt_text"] == "[3 tokens]" + assert result["completion_text"] == "[0 tokens]" + assert result["segments"] is None + + def test_all_ones_gen_only(self, tokenizer): + ids = [10, 20, 30] + mask = [1, 1, 1] + result = WorkflowExecutor._split_trajectory_for_dump(ids, mask, tokenizer) + assert result["prompt_end"] == 0 + assert result["prompt_text"] == "[0 tokens]" + assert result["completion_text"] == "[3 tokens]" + assert result["segments"] is None + + def test_three_gen_runs(self, tokenizer): + ids = list(range(12)) + mask = [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1] + result = WorkflowExecutor._split_trajectory_for_dump(ids, mask, tokenizer) + assert result["prompt_end"] == 2 + assert result["segments"] is not None + roles = [s["role"] for s in result["segments"]] + assert roles == ["prompt", "gen", "context", "gen", "context", "gen"] + + def test_length_mismatch_raises_valueerror(self, tokenizer): + with pytest.raises(ValueError, match="ids length 2 != mask length 1"): + WorkflowExecutor._split_trajectory_for_dump([1, 2], [0], tokenizer) + + def test_single_token_gen(self, tokenizer): + ids = [42] + mask = [1] + result = WorkflowExecutor._split_trajectory_for_dump(ids, mask, tokenizer) + assert result["prompt_end"] == 0 + assert result["segments"] is None + + +class TestComputeOutputVersions: + def test_filters_negative_one_placeholders(self): + versions = [-1, -1, 5, 5, 6] + mask = [0, 0, 1, 1, 1] + head, tail, rle = WorkflowExecutor._compute_output_versions(versions, mask) + assert head == 5 + assert tail == 6 + assert rle == [[5, 2], [6, 1]] + + def test_single_version(self): + versions = [-1, 3, 3, 3] + mask = [0, 1, 1, 1] + head, tail, rle = WorkflowExecutor._compute_output_versions(versions, mask) + assert head == 3 + assert tail == 3 + assert rle == [[3, 3]] + + def test_multiple_version_transitions(self): + versions = [-1, 2, 2, 3, 3, 4] + mask = [0, 1, 1, 1, 1, 1] + head, tail, rle = WorkflowExecutor._compute_output_versions(versions, mask) + assert head == 2 + assert tail == 4 + assert rle == [[2, 2], [3, 2], [4, 1]] + + def test_all_masked_out(self): + versions = [1, 2, 3] + mask = [0, 0, 0] + head, tail, rle = WorkflowExecutor._compute_output_versions(versions, mask) + assert head == -1 + assert tail == -1 + assert rle == [] + + def test_interleaved_multi_turn(self): + versions = [-1, -1, 5, 5, -1, -1, 6, 6] + mask = [0, 0, 1, 1, 0, 0, 1, 1] + head, tail, rle = WorkflowExecutor._compute_output_versions(versions, mask) + assert head == 5 + assert tail == 6 + assert rle == [[5, 2], [6, 2]] diff --git a/tests/test_stats_tracker_reduce_group.py b/tests/test_stats_tracker_reduce_group.py new file mode 100644 index 0000000000..f7b04b8c73 --- /dev/null +++ b/tests/test_stats_tracker_reduce_group.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 + +from unittest.mock import patch + +import torch + +from areal.utils.stats_tracker import DistributedStatsTracker + + +def test_stats_tracker_scalar_reduce_group_override_used_on_export(): + # Arrange + tracker = DistributedStatsTracker() + default_group = object() + override_group = object() + + tracker.scalar(score=1.0, reduce_group=override_group) + + # Act + with patch("areal.utils.stats_tracker.dist.all_reduce") as mock_all_reduce: + tracker.export(key="score", reduce_group=default_group) + + # Assert + assert mock_all_reduce.call_count == 2 + assert all( + call.kwargs["group"] is override_group for call in mock_all_reduce.mock_calls + ) + + +def test_stats_tracker_stat_reduce_group_override_used_on_export(): + # Arrange + tracker = DistributedStatsTracker() + default_group = object() + override_group = object() + + tracker.denominator(mask=torch.tensor([True, False, True])) + tracker.stat( + "mask", + loss=torch.tensor([1.0, 2.0, 3.0]), + reduce_group=override_group, + ) + + # Act + with patch("areal.utils.stats_tracker.dist.all_reduce") as mock_all_reduce: + tracker.export(key="loss", reduce_group=default_group) + + # Assert + assert mock_all_reduce.call_count == 4 + assert all( + call.kwargs["group"] is override_group for call in mock_all_reduce.mock_calls + ) diff --git a/tests/test_stepwise_reward_stage1.py b/tests/test_stepwise_reward_stage1.py new file mode 100644 index 0000000000..381ba6820d --- /dev/null +++ b/tests/test_stepwise_reward_stage1.py @@ -0,0 +1,215 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import contextlib +import sys +import types + +import pytest +import torch + +from areal.api.reward_api import normalize_reward_result +from areal.api.reward_api import RewardResult + + +@contextlib.asynccontextmanager +async def _dummy_atrace_session_phase(_name): + yield + + +def _dummy_trace_session(_name): + def decorator(fn): + return fn + + return decorator + + +def _dummy_session_context(): + def decorator(fn): + return fn + + return decorator + + +_perf_tracer_stub = types.ModuleType("areal.utils.perf_tracer") +_perf_tracer_stub.atrace_session_phase = _dummy_atrace_session_phase +_perf_tracer_stub.trace_session = _dummy_trace_session +_perf_tracer_stub.session_context = _dummy_session_context +sys.modules.setdefault("areal.utils.perf_tracer", _perf_tracer_stub) + + +class _DummyTracker: + def scalar(self, **kwargs): + return None + + +_stats_tracker_stub = types.ModuleType("areal.utils.stats_tracker") +_stats_tracker_stub.get = lambda _scope: _DummyTracker() +sys.modules.setdefault("areal.utils.stats_tracker", _stats_tracker_stub) + +_workflow_context_stub = types.ModuleType("areal.infra.workflow_context") +_workflow_context_stub.stat_scope = lambda: "rollout" +sys.modules.setdefault("areal.infra.workflow_context", _workflow_context_stub) + +_infra_stub = types.ModuleType("areal.infra") +_infra_stub.workflow_context = _workflow_context_stub +sys.modules.setdefault("areal.infra", _infra_stub) + +from areal.workflow.rlvr import RLVRWorkflow + + +class _DummyTokenizer: + eos_token_id = 0 + pad_token_id = 0 + + def decode(self, token_ids): + return "|".join(str(x) for x in token_ids) + + +class _DummyGConfig: + def __init__(self, max_new_tokens=8, n_samples=1): + self.max_new_tokens = max_new_tokens + self.n_samples = n_samples + + def new_with_stop_and_pad_token_ids(self, _tokenizer): + return self + + def new(self, n_samples=1): + return _DummyGConfig( + max_new_tokens=self.max_new_tokens, + n_samples=n_samples, + ) + + +class _DummyModelResponse: + def __init__( + self, + input_tokens, + output_tokens, + output_logprobs, + output_versions, + stop_reason="stop", + tokenizer=None, + ): + self.input_tokens = input_tokens + self.output_tokens = output_tokens + self.output_logprobs = output_logprobs + self.output_versions = output_versions + self.stop_reason = stop_reason + self.tokenizer = tokenizer + + @property + def input_len(self): + return len(self.input_tokens) + + @property + def output_len(self): + return len(self.output_tokens) + + +class _DummyEngine: + async def agenerate(self, req): + return _DummyModelResponse( + input_tokens=req.input_ids, + output_tokens=[11, 12, 13, 14], + output_logprobs=[-0.1, -0.2, -0.3, -0.4], + output_versions=[1, 1, 1, 1], + stop_reason="stop", + tokenizer=req.tokenizer, + ) + + +def _dummy_get_input_ids_fn(_data, _tokenizer, _enable_thinking): + return [101, 102] + + +def _scalar_reward_fn(prompt, completions, prompt_ids, completion_ids, **kwargs): + del prompt, completions, prompt_ids, completion_ids, kwargs + return 1.25 + + +def _stepwise_reward_fn(prompt, completions, prompt_ids, completion_ids, **kwargs): + del prompt, completions, prompt_ids, completion_ids, kwargs + return RewardResult( + final_reward=0.5, + step_rewards=[0.2, -0.1], + step_ends=[2, 4], + metadata={"num_steps": 2}, + ) + + +def _bad_reward_fn(prompt, completions, prompt_ids, completion_ids, **kwargs): + del prompt, completions, prompt_ids, completion_ids, kwargs + return RewardResult( + final_reward=0.0, + step_rewards=[1.0], + step_ends=[5], + ) + + +class TestRewardResultNormalization: + def test_normalize_scalar_reward(self): + reward = normalize_reward_result(1.5) + assert reward == RewardResult(final_reward=1.5) + + def test_normalize_reward_result_passthrough(self): + reward = RewardResult(final_reward=2.0, step_rewards=[0.1], step_ends=[1]) + assert normalize_reward_result(reward) is reward + + +class TestRLVRWorkflowStepwiseReward: + @pytest.mark.asyncio + async def test_scalar_reward_keeps_zero_stepwise_tensors(self): + workflow = RLVRWorkflow( + reward_fn=_scalar_reward_fn, + gconfig=_DummyGConfig(max_new_tokens=8), + tokenizer=_DummyTokenizer(), + enable_thinking=False, + get_input_ids_fn=_dummy_get_input_ids_fn, + ) + + result = await workflow.arun_episode(_DummyEngine(), {"messages": []}) + + assert torch.equal(result["rewards"], torch.tensor([1.25], dtype=torch.float32)) + assert torch.equal( + result["step_rewards"], torch.zeros((1, 6), dtype=torch.float32) + ) + assert torch.equal( + result["step_reward_mask"], torch.zeros((1, 6), dtype=torch.bool) + ) + + @pytest.mark.asyncio + async def test_stepwise_reward_aligns_to_completion_end_positions(self): + workflow = RLVRWorkflow( + reward_fn=_stepwise_reward_fn, + gconfig=_DummyGConfig(max_new_tokens=8), + tokenizer=_DummyTokenizer(), + enable_thinking=False, + get_input_ids_fn=_dummy_get_input_ids_fn, + ) + + result = await workflow.arun_episode(_DummyEngine(), {"messages": []}) + + assert torch.equal(result["rewards"], torch.tensor([0.5], dtype=torch.float32)) + expected_step_rewards = torch.tensor( + [[0.0, 0.0, 0.0, 0.2, 0.0, -0.1]], dtype=torch.float32 + ) + expected_mask = torch.tensor( + [[False, False, False, True, False, True]], dtype=torch.bool + ) + assert torch.equal(result["step_rewards"], expected_step_rewards) + assert torch.equal(result["step_reward_mask"], expected_mask) + + @pytest.mark.asyncio + async def test_invalid_stepwise_reward_raises(self): + workflow = RLVRWorkflow( + reward_fn=_bad_reward_fn, + gconfig=_DummyGConfig(max_new_tokens=8), + tokenizer=_DummyTokenizer(), + enable_thinking=False, + get_input_ids_fn=_dummy_get_input_ids_fn, + ) + + with pytest.raises(ValueError, match="Invalid step_end"): + await workflow.arun_episode(_DummyEngine(), {"messages": []}) diff --git a/tests/test_train_engine.py b/tests/test_train_engine.py index d5090b2fe7..3a72fe4f0d 100644 --- a/tests/test_train_engine.py +++ b/tests/test_train_engine.py @@ -325,6 +325,7 @@ def from_pretrained(pretrained_model_name_or_path, **kwargs): trial_name="trial0", path="test-model", attn_impl="kernels-community/flash-attn", + optimizer_dtype="bfloat16", fsdp=FSDPEngineConfig(memory_efficient_load=memory_efficient_load), ) engine = fsdp_module.FSDPEngine(config) diff --git a/tests/test_vllm_generation_request.py b/tests/test_vllm_generation_request.py new file mode 100644 index 0000000000..18a42d0ee8 --- /dev/null +++ b/tests/test_vllm_generation_request.py @@ -0,0 +1,20 @@ +from areal.api.cli_args import GenerationHyperparameters +from areal.api.io_struct import ModelRequest +from areal.engine.vllm_remote import VLLMBackend + + +def test_vllm_forwards_frequency_penalty_and_stop(): + """The vLLM backend must forward frequency_penalty and stop like the SGLang + backend does; both are GenerationHyperparameters and are accepted by vLLM's + OpenAI-compatible /v1/completions endpoint.""" + gconfig = GenerationHyperparameters( + max_new_tokens=8, frequency_penalty=0.5, stop=["STOP"] + ) + req = ModelRequest(input_ids=[11, 12], gconfig=gconfig) + + payload = ( + VLLMBackend().build_generation_request(req, with_lora=False, version=0).payload + ) + + assert payload["frequency_penalty"] == 0.5 + assert payload["stop"] == ["STOP"] diff --git a/tests/v2/__init__.py b/tests/v2/__init__.py new file mode 100644 index 0000000000..a72302f3f2 --- /dev/null +++ b/tests/v2/__init__.py @@ -0,0 +1 @@ +"""V2 tests.""" diff --git a/tests/experimental/agent_service/__init__.py b/tests/v2/agent_service/__init__.py similarity index 100% rename from tests/experimental/agent_service/__init__.py rename to tests/v2/agent_service/__init__.py diff --git a/tests/experimental/agent_service/test_agent_router.py b/tests/v2/agent_service/test_agent_router.py similarity index 94% rename from tests/experimental/agent_service/test_agent_router.py rename to tests/v2/agent_service/test_agent_router.py index 13e682de23..365859bdd2 100644 --- a/tests/experimental/agent_service/test_agent_router.py +++ b/tests/v2/agent_service/test_agent_router.py @@ -4,9 +4,9 @@ import pytest -from areal.experimental.agent_service.auth import DEFAULT_ADMIN_API_KEY, admin_headers -from areal.experimental.agent_service.router.app import create_router_app -from areal.experimental.agent_service.router.config import RouterConfig +from areal.v2.agent_service.auth import DEFAULT_ADMIN_API_KEY, admin_headers +from areal.v2.agent_service.router.app import create_router_app +from areal.v2.agent_service.router.config import RouterConfig httpx = pytest.importorskip("httpx") diff --git a/tests/experimental/agent_service/test_agent_worker.py b/tests/v2/agent_service/test_agent_worker.py similarity index 95% rename from tests/experimental/agent_service/test_agent_worker.py rename to tests/v2/agent_service/test_agent_worker.py index 740622f877..4bcaa1d52a 100644 --- a/tests/experimental/agent_service/test_agent_worker.py +++ b/tests/v2/agent_service/test_agent_worker.py @@ -6,13 +6,13 @@ import pytest -from areal.experimental.agent_service.types import ( +from areal.v2.agent_service.types import ( AgentRequest, AgentResponse, AgentRunnable, EventEmitter, ) -from areal.experimental.agent_service.worker.app import ( +from areal.v2.agent_service.worker.app import ( _CollectingEmitter, create_worker_app, ) @@ -50,7 +50,7 @@ async def run( def _make_client(agent_cls): with patch( - "areal.experimental.agent_service.worker.app.import_from_string", + "areal.v2.agent_service.worker.app.import_from_string", return_value=agent_cls, ): app = create_worker_app("mock.path") diff --git a/tests/experimental/agent_service/test_controller.py b/tests/v2/agent_service/test_controller.py similarity index 97% rename from tests/experimental/agent_service/test_controller.py rename to tests/v2/agent_service/test_controller.py index 6f6aeb5f49..a7f846dfa9 100644 --- a/tests/experimental/agent_service/test_controller.py +++ b/tests/v2/agent_service/test_controller.py @@ -16,9 +16,9 @@ import pytest from areal.api.cli_args import AgentConfig, SchedulingSpec -from areal.experimental.agent_service.controller.controller import AgentController +from areal.v2.agent_service.controller.controller import AgentController -CTRL = "areal.experimental.agent_service.controller.controller" +CTRL = "areal.v2.agent_service.controller.controller" @dataclass @@ -170,10 +170,7 @@ def mock_post(url, **kwargs): create_call = scheduler.create_workers.call_args job = create_call.kwargs.get("job") or create_call.args[0] assert job.tasks[0].env_vars == {"ANTHROPIC_API_KEY": "test-anthropic-key"} - assert ( - job.tasks[0].cmd - == f"{sys.executable} -m areal.experimental.agent_service.guard" - ) + assert job.tasks[0].cmd == f"{sys.executable} -m areal.v2.agent_service.guard" assert all( payload.get("env") == {"ANTHROPIC_API_KEY": "test-anthropic-key"} for payload in fork_payloads diff --git a/tests/experimental/agent_service/test_guard.py b/tests/v2/agent_service/test_guard.py similarity index 91% rename from tests/experimental/agent_service/test_guard.py rename to tests/v2/agent_service/test_guard.py index 674808d618..e9e7dfdbfc 100644 --- a/tests/experimental/agent_service/test_guard.py +++ b/tests/v2/agent_service/test_guard.py @@ -4,7 +4,7 @@ The agent_blueprint has been removed in v2 — all orchestration logic now lives in AgentController. -Test structure mirrors ``tests/experimental/inference_service/test_guard.py``. +Test structure mirrors ``tests/v2/inference_service/test_guard.py``. """ from __future__ import annotations @@ -13,8 +13,8 @@ import pytest -from areal.experimental.agent_service.guard import app as guard_module -from areal.experimental.agent_service.guard.app import app +from areal.v2.agent_service.guard import app as guard_module +from areal.v2.agent_service.guard.app import app GUARD_APP = "areal.infra.rpc.guard.app" diff --git a/tests/v2/agent_service/test_integration.py b/tests/v2/agent_service/test_integration.py new file mode 100644 index 0000000000..9735cfc85b --- /dev/null +++ b/tests/v2/agent_service/test_integration.py @@ -0,0 +1,900 @@ +"""Integration tests for the Agent Service. + +Tests the full HTTP microservice stack: Worker → DataProxy → Router, +plus utility functions from the Bridge and Gateway health endpoints. +""" + +from __future__ import annotations + +import json +from unittest.mock import patch + +import pytest + +from areal.v2.agent_service.auth import DEFAULT_ADMIN_API_KEY, admin_headers +from areal.v2.agent_service.data_proxy.app import create_data_proxy_app +from areal.v2.agent_service.data_proxy.config import DataProxyConfig +from areal.v2.agent_service.gateway.app import create_gateway_app +from areal.v2.agent_service.gateway.bridge import ( + SESSION_KEY_HEADER, + ChatCompletionsBridge, + OpenResponsesBridge, +) +from areal.v2.agent_service.gateway.config import GatewayConfig +from areal.v2.agent_service.protocol import PASSTHROUGH_HEADER +from areal.v2.agent_service.router.app import create_router_app +from areal.v2.agent_service.router.config import RouterConfig +from areal.v2.agent_service.types import ( + AgentRequest, + AgentResponse, + EventEmitter, + StreamResponse, +) +from areal.v2.agent_service.worker.app import create_worker_app + +httpx = pytest.importorskip("httpx") + +_AUTH = admin_headers(DEFAULT_ADMIN_API_KEY) + + +class _EchoAgent: + async def run( + self, request: AgentRequest, *, emitter: EventEmitter + ) -> AgentResponse: + history_summary = f"history={len(request.history)}" + await emitter.emit_delta(f"echo: {request.message} ({history_summary})") + return AgentResponse(summary=f"echo: {request.message}") + + +class _ToolAgent: + async def run( + self, request: AgentRequest, *, emitter: EventEmitter + ) -> AgentResponse: + await emitter.emit_tool_call("lookup", '{"id": "123"}') + await emitter.emit_tool_result("lookup", '{"status": "ok"}') + await emitter.emit_delta("Lookup complete") + return AgentResponse( + summary="Lookup complete", + metadata={"tool_calls": [{"name": "lookup", "arguments": {"id": "123"}}]}, + ) + + +class _StreamAgent: + """Returns a :class:`StreamResponse` from ``run`` when the turn carries a + ``chat_request`` (the raw-passthrough channel), echoing the request back as + an SSE body so tests can assert the worker relays it verbatim. Otherwise it + falls back to a structured :class:`AgentResponse`.""" + + async def run( + self, request: AgentRequest, *, emitter: EventEmitter + ) -> AgentResponse | StreamResponse: + chat_request = (request.metadata or {}).get("chat_request") + if chat_request is None: + return AgentResponse(summary="structured") + + async def _body(): + yield f"data: msg={request.message}\n\n".encode() + yield f"data: chat_request={json.dumps(chat_request)}\n\n".encode() + yield b"data: [DONE]\n\n" + + return StreamResponse( + status_code=200, + headers={"content-type": "text/event-stream"}, + body=_body(), + ) + + +class _JsonPassthroughAgent: + """Returns a non-streaming :class:`StreamResponse` whose body is + ``application/json``. Used to prove the DataProxy keys its relay-vs-parse + decision on the passthrough marker, not on ``Content-Type``: without the + marker this body would be mistaken for a structured turn and parsed.""" + + async def run( + self, request: AgentRequest, *, emitter: EventEmitter + ) -> StreamResponse: + async def _body(): + yield json.dumps({"verbatim": request.message, "events": []}).encode() + + return StreamResponse( + status_code=200, + headers={"content-type": "application/json"}, + body=_body(), + ) + + +def _make_worker_app(agent_cls): + with patch( + "areal.v2.agent_service.worker.app.import_from_string", + return_value=agent_cls, + ): + return create_worker_app("mock.path") + + +def _worker_patches(worker_transport): + """Route the DataProxy's worker calls through the worker ASGI app. + + The DataProxy opens the worker ``/run`` stream via ``client.send`` (so it + can sniff the response ``Content-Type`` before reading) and closes worker + sessions via ``client.post``; patch both to forward to ``worker_transport``. + """ + original_post = httpx.AsyncClient.post + original_send = httpx.AsyncClient.send + + async def _forward(path, body): + # Use the unbound original ``send`` so the inner ASGI call does not + # re-enter the patched ``post``/``send`` (which would recurse forever). + async with httpx.AsyncClient( + transport=worker_transport, base_url="http://worker" + ) as wc: + req = wc.build_request("POST", path, json=body) + return await original_send(wc, req) + + async def patched_post(self, url, **kwargs): + if "http://worker" in str(url): + path = str(url).split("http://worker")[-1] + return await _forward(path, kwargs.get("json")) + return await original_post(self, url, **kwargs) + + async def patched_send(self, request, **kwargs): + if "http://worker" in str(request.url): + path = str(request.url).split("http://worker")[-1] + body = json.loads(request.content) if request.content else None + r = await _forward(path, body) + headers = { + "content-type": r.headers.get("content-type", "application/json") + } + # Preserve the passthrough marker so the DataProxy's relay-vs-parse + # decision is exercised faithfully end-to-end. + if PASSTHROUGH_HEADER in r.headers: + headers[PASSTHROUGH_HEADER] = r.headers[PASSTHROUGH_HEADER] + # Return the body as a stream (the real DataProxy opens the worker + # with ``stream=True``); a bytes payload would refuse ``aiter_raw``. + raw = r.content + + async def _stream(): + yield raw + + return httpx.Response( + r.status_code, + headers=headers, + content=_stream(), + request=request, + ) + return await original_send(self, request, **kwargs) + + return patched_post, patched_send + + +class TestWorkerDataProxyIntegration: + """Test DataProxy → Worker chain using ASGITransport for the Worker.""" + + @pytest.mark.asyncio + async def test_single_turn(self): + worker_app = _make_worker_app(_EchoAgent) + worker_transport = httpx.ASGITransport(app=worker_app) + + async with httpx.AsyncClient( + transport=worker_transport, base_url="http://worker" + ) as worker_client: + # DataProxy forwards to worker — test worker directly first + resp = await worker_client.post( + "/run", + json={ + "message": "hello", + "session_key": "s1", + "run_id": "r1", + "history": [], + }, + ) + assert resp.status_code == 200 + data = resp.json() + assert "echo: hello" in data["summary"] + + @pytest.mark.asyncio + async def test_data_proxy_manages_history(self): + worker_app = _make_worker_app(_EchoAgent) + worker_transport = httpx.ASGITransport(app=worker_app) + + proxy_app = create_data_proxy_app(DataProxyConfig(worker_addr="http://worker")) + patched_post, patched_send = _worker_patches(worker_transport) + proxy_transport = httpx.ASGITransport(app=proxy_app) + + with ( + patch.object(httpx.AsyncClient, "post", patched_post), + patch.object(httpx.AsyncClient, "send", patched_send), + ): + async with httpx.AsyncClient( + transport=proxy_transport, base_url="http://proxy" + ) as proxy_client: + # Turn 1 + r1 = await proxy_client.post( + "/session/s1/turn", + json={"message": "hello", "run_id": "r1"}, + ) + assert r1.status_code == 200 + assert "echo: hello" in r1.json()["summary"] + + # Turn 2 — history should have turn 1 + r2 = await proxy_client.post( + "/session/s1/turn", + json={"message": "world", "run_id": "r2"}, + ) + assert r2.status_code == 200 + + # Check history grew + h = await proxy_client.get("/session/s1/history") + history = h.json()["history"] + assert len(history) >= 2 # at least user+assistant from turn 1 + + @pytest.mark.asyncio + async def test_close_session_clears_history(self): + worker_app = _make_worker_app(_EchoAgent) + worker_transport = httpx.ASGITransport(app=worker_app) + proxy_app = create_data_proxy_app(DataProxyConfig(worker_addr="http://worker")) + patched_post, patched_send = _worker_patches(worker_transport) + proxy_transport = httpx.ASGITransport(app=proxy_app) + + with ( + patch.object(httpx.AsyncClient, "post", patched_post), + patch.object(httpx.AsyncClient, "send", patched_send), + ): + async with httpx.AsyncClient( + transport=proxy_transport, base_url="http://proxy" + ) as proxy_client: + await proxy_client.post( + "/session/s1/turn", + json={"message": "hi", "run_id": "r1"}, + ) + await proxy_client.post("/session/s1/close") + h = await proxy_client.get("/session/s1/history") + assert h.json()["history"] == [] + + @pytest.mark.asyncio + async def test_json_passthrough_relayed_verbatim_not_parsed(self): + """A non-streaming passthrough body is itself ``application/json``; the + DataProxy must relay it verbatim on the marker, never mistaking it for a + structured turn (which would rebuild history and re-serialise).""" + worker_app = _make_worker_app(_JsonPassthroughAgent) + worker_transport = httpx.ASGITransport(app=worker_app) + proxy_app = create_data_proxy_app(DataProxyConfig(worker_addr="http://worker")) + patched_post, patched_send = _worker_patches(worker_transport) + proxy_transport = httpx.ASGITransport(app=proxy_app) + + with ( + patch.object(httpx.AsyncClient, "post", patched_post), + patch.object(httpx.AsyncClient, "send", patched_send), + ): + async with httpx.AsyncClient( + transport=proxy_transport, base_url="http://proxy" + ) as proxy_client: + r = await proxy_client.post( + "/session/s1/turn", + json={"message": "hi", "run_id": "r1"}, + ) + assert r.status_code == 200 + # Body relayed byte-for-byte (the agent's exact JSON). + assert r.json() == {"verbatim": "hi", "events": []} + # No history was kept on the passthrough path. + h = await proxy_client.get("/session/s1/history") + assert h.json()["history"] == [] + + +class TestRouterIntegration: + @pytest.mark.asyncio + async def test_register_and_route(self): + router_app = create_router_app( + RouterConfig(admin_api_key=DEFAULT_ADMIN_API_KEY) + ) + transport = httpx.ASGITransport(app=router_app) + + async with httpx.AsyncClient( + transport=transport, base_url="http://router" + ) as client: + await client.post( + "/register", + json={"addr": "http://proxy1:9100"}, + headers=_AUTH, + ) + resp = await client.post( + "/route", json={"session_key": "s1"}, headers=_AUTH + ) + assert resp.json()["data_proxy_addr"] == "http://proxy1:9100" + + resp2 = await client.post( + "/route", json={"session_key": "s1"}, headers=_AUTH + ) + assert resp2.json()["data_proxy_addr"] == "http://proxy1:9100" + + +class TestToolCallFlow: + @pytest.mark.asyncio + async def test_tool_events_through_proxy(self): + worker_app = _make_worker_app(_ToolAgent) + worker_transport = httpx.ASGITransport(app=worker_app) + proxy_app = create_data_proxy_app(DataProxyConfig(worker_addr="http://worker")) + patched_post, patched_send = _worker_patches(worker_transport) + proxy_transport = httpx.ASGITransport(app=proxy_app) + + with ( + patch.object(httpx.AsyncClient, "post", patched_post), + patch.object(httpx.AsyncClient, "send", patched_send), + ): + async with httpx.AsyncClient( + transport=proxy_transport, base_url="http://proxy" + ) as proxy_client: + resp = await proxy_client.post( + "/session/s1/turn", + json={"message": "lookup 123", "run_id": "r1"}, + ) + data = resp.json() + assert data["summary"] == "Lookup complete" + events = data["events"] + types = {e["type"] for e in events} + assert "tool_call" in types + assert "tool_result" in types + + # History should include tool call records + h = await proxy_client.get("/session/s1/history") + history = h.json()["history"] + tool_msgs = [m for m in history if m.get("role") == "tool"] + assert len(tool_msgs) > 0 + assert "tool_call_id" in tool_msgs[0] + + +class TestArealInference: + """Self-evolution (decoupled from the training side): the **caller** mints + its own per-session ``sk-sess-*`` and passes it on the turn body as + ``session_api_key``. The DataProxy never contacts ``/rl/start_session`` — + it caches the routing handle on the session and injects ``areal_inference`` + into the worker request metadata. A multi-turn session may send the fields + on the first turn and omit them afterwards. Without the flag the worker + metadata stays clean. + """ + + @staticmethod + def _patched_send(worker_bodies, *, stream=False): + original_send = httpx.AsyncClient.send + + async def patched_send(self, request, **kwargs): + if str(request.url).endswith("/run"): + worker_bodies.append(json.loads(request.content)) + if stream: + + async def _s(): + yield b"data: ok\n\n" + + return httpx.Response( + 200, + headers={ + "content-type": "text/event-stream", + PASSTHROUGH_HEADER: "1", + }, + content=_s(), + request=request, + ) + return httpx.Response( + 200, + json={"summary": "ok", "events": [], "metadata": {}}, + request=request, + ) + return await original_send(self, request, **kwargs) + + return patched_send + + @pytest.mark.asyncio + async def test_turn_caches_inference_routing_once(self): + proxy_app = create_data_proxy_app(DataProxyConfig(worker_addr="http://worker")) + worker_bodies: list[dict] = [] + proxy_transport = httpx.ASGITransport(app=proxy_app) + + with patch.object(httpx.AsyncClient, "send", self._patched_send(worker_bodies)): + async with httpx.AsyncClient( + transport=proxy_transport, base_url="http://proxy" + ) as client: + turn_body = { + "message": "hi", + "run_id": "r1", + "inf_base_url": "http://inf-gateway", + "session_api_key": "sk-sess-xyz", + "inf_model": "Qwen", + } + r1 = await client.post("/session/s1/turn", json=turn_body) + assert r1.status_code == 200 + # The second turn omits the inference fields; the proxy reuses + # the handle cached on the session. + r2 = await client.post( + "/session/s1/turn", json={"message": "again", "run_id": "r2"} + ) + assert r2.status_code == 200 + + # The worker received the inference routing handle on both turns. + assert len(worker_bodies) == 2 + for body in worker_bodies: + inf = body["metadata"]["areal_inference"] + assert inf["base_url"] == "http://inf-gateway" + assert inf["api_key"] == "sk-sess-xyz" + assert inf["model"] == "Qwen" + + # The Agent Service never surfaces a training-side session id back. + assert "areal_inference" not in r2.json().get("metadata", {}) + + @pytest.mark.asyncio + async def test_turn_without_inference_fields_skips_inference(self): + proxy_app = create_data_proxy_app(DataProxyConfig(worker_addr="http://worker")) + worker_bodies: list[dict] = [] + proxy_transport = httpx.ASGITransport(app=proxy_app) + + with patch.object(httpx.AsyncClient, "send", self._patched_send(worker_bodies)): + async with httpx.AsyncClient( + transport=proxy_transport, base_url="http://proxy" + ) as client: + r = await client.post( + "/session/s2/turn", json={"message": "hi", "run_id": "r1"} + ) + assert r.status_code == 200 + + assert len(worker_bodies) == 1 + assert "areal_inference" not in worker_bodies[0]["metadata"] + + @pytest.mark.asyncio + async def test_turn_requires_both_inference_fields(self): + proxy_app = create_data_proxy_app(DataProxyConfig(worker_addr="http://worker")) + worker_bodies: list[dict] = [] + proxy_transport = httpx.ASGITransport(app=proxy_app) + + with patch.object(httpx.AsyncClient, "send", self._patched_send(worker_bodies)): + async with httpx.AsyncClient( + transport=proxy_transport, base_url="http://proxy" + ) as client: + # inf_base_url present but session_api_key missing → 400 before + # any worker call (the pair is required once either is sent). + r = await client.post( + "/session/s3/turn", + json={"message": "hi", "inf_base_url": "http://inf-gateway"}, + ) + assert r.status_code == 400 + + assert worker_bodies == [] + + @pytest.mark.asyncio + async def test_stream_turn_injects_inference(self): + """The same routing injection applies to the raw-passthrough channel: + when the worker returns a stream, the proxy still injects + ``areal_inference`` into the worker request and relays the body.""" + proxy_app = create_data_proxy_app(DataProxyConfig(worker_addr="http://worker")) + worker_bodies: list[dict] = [] + proxy_transport = httpx.ASGITransport(app=proxy_app) + + with patch.object( + httpx.AsyncClient, "send", self._patched_send(worker_bodies, stream=True) + ): + async with httpx.AsyncClient( + transport=proxy_transport, base_url="http://proxy" + ) as client: + r = await client.post( + "/session/s1/turn", + json={ + "message": "hi", + "run_id": "r1", + "inf_base_url": "http://inf-gateway", + "session_api_key": "sk-sess-xyz", + "inf_model": "Qwen", + }, + ) + assert r.status_code == 200 + assert r.text == "data: ok\n\n" + + assert len(worker_bodies) == 1 + inf = worker_bodies[0]["metadata"]["areal_inference"] + assert inf["base_url"] == "http://inf-gateway" + assert inf["api_key"] == "sk-sess-xyz" + assert inf["model"] == "Qwen" + + +class TestWorkerRunStream: + """The single ``/run`` endpoint relays an agent's :class:`StreamResponse` + verbatim (tagged with the passthrough marker header so the DataProxy relays + rather than parses), and serialises an :class:`AgentResponse` to JSON + otherwise.""" + + @pytest.mark.asyncio + async def test_run_relays_stream_verbatim_with_metadata(self): + worker_app = _make_worker_app(_StreamAgent) + worker_transport = httpx.ASGITransport(app=worker_app) + + async with httpx.AsyncClient( + transport=worker_transport, base_url="http://worker" + ) as client: + async with client.stream( + "POST", + "/run", + json={ + "message": "hi", + "session_key": "s1", + "run_id": "r1", + "history": [], + "metadata": {"chat_request": {"model": "m", "stream": True}}, + }, + ) as resp: + assert resp.status_code == 200 + assert resp.headers["content-type"].startswith("text/event-stream") + # The worker tags raw-passthrough turns for the DataProxy. + assert resp.headers.get(PASSTHROUGH_HEADER) == "1" + body = b"" + async for chunk in resp.aiter_bytes(): + body += chunk + + text = body.decode() + assert "data: msg=hi\n\n" in text + assert 'data: chat_request={"model": "m", "stream": true}\n\n' in text + assert text.endswith("data: [DONE]\n\n") + + @pytest.mark.asyncio + async def test_run_returns_json_for_structured_response(self): + worker_app = _make_worker_app(_StreamAgent) + worker_transport = httpx.ASGITransport(app=worker_app) + + async with httpx.AsyncClient( + transport=worker_transport, base_url="http://worker" + ) as client: + resp = await client.post( + "/run", + json={"message": "hi", "session_key": "s1", "run_id": "r1"}, + ) + assert resp.status_code == 200 + assert resp.headers["content-type"].startswith("application/json") + # No marker on a structured turn → the DataProxy parses it. + assert PASSTHROUGH_HEADER not in resp.headers + assert resp.json()["summary"] == "structured" + + +class TestGatewayHealth: + @pytest.mark.asyncio + async def test_health(self): + app = create_gateway_app(GatewayConfig(router_addr="http://fake-router")) + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient( + transport=transport, base_url="http://gw" + ) as client: + resp = await client.get("/health") + assert resp.status_code == 200 + assert resp.json()["status"] == "ok" + + +class TestGatewayCloseSession: + """The gateway's ``POST /sessions/close`` (session_key in the body) resolves + the owning DataProxy via the Router and forwards the close, mirroring the + internal ``/session/{key}/close`` endpoint shape.""" + + @staticmethod + def _patched_post(calls): + original_post = httpx.AsyncClient.post + + async def patched_post(self, url, **kwargs): + url = str(url) + if url.endswith("/route"): + return httpx.Response( + 200, + json={"data_proxy_addr": "http://proxy1"}, + request=httpx.Request("POST", url), + ) + if url.startswith("http://proxy1") and url.endswith("/close"): + calls.append(url) + return httpx.Response( + 200, json={"status": "ok"}, request=httpx.Request("POST", url) + ) + return await original_post(self, url, **kwargs) + + return patched_post + + @pytest.mark.asyncio + async def test_close_forwards_to_data_proxy(self): + app = create_gateway_app(GatewayConfig(router_addr="http://fake-router")) + transport = httpx.ASGITransport(app=app) + calls: list[str] = [] + + with patch.object(httpx.AsyncClient, "post", self._patched_post(calls)): + async with httpx.AsyncClient( + transport=transport, base_url="http://gw" + ) as client: + resp = await client.post( + "/sessions/close", json={"session_key": "s1"}, headers=_AUTH + ) + assert resp.status_code == 200 + assert resp.json()["status"] == "ok" + + assert calls == ["http://proxy1/session/s1/close"] + + @pytest.mark.asyncio + async def test_close_requires_session_key(self): + app = create_gateway_app(GatewayConfig(router_addr="http://fake-router")) + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient( + transport=transport, base_url="http://gw" + ) as client: + resp = await client.post("/sessions/close", json={}, headers=_AUTH) + assert resp.status_code == 400 + + @pytest.mark.asyncio + async def test_close_requires_admin_key(self): + app = create_gateway_app(GatewayConfig(router_addr="http://fake-router")) + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient( + transport=transport, base_url="http://gw" + ) as client: + resp = await client.post("/sessions/close", json={"session_key": "s1"}) + assert resp.status_code in (401, 422) + + +class TestBridgeExtractMessage: + def test_text_message(self): + items = [ + { + "type": "message", + "content": [{"type": "input_text", "text": "Hello"}], + } + ] + assert OpenResponsesBridge._extract_message(items, "") == "Hello" + + def test_string_content(self): + items = [{"type": "message", "content": "Simple"}] + assert OpenResponsesBridge._extract_message(items, "") == "Simple" + + def test_instructions_prepended(self): + items = [{"type": "message", "content": "Hi"}] + result = OpenResponsesBridge._extract_message(items, "Be helpful") + assert result.startswith("Be helpful") + assert "Hi" in result + + def test_function_call_output(self): + items = [{"type": "function_call_output", "output": "42"}] + result = OpenResponsesBridge._extract_message(items, "") + assert "[tool result] 42" in result + + +class TestBridgeDeriveSessionKey: + def test_with_user(self): + key = OpenResponsesBridge._derive_session_key("user1", "model1") + assert key == "agent:model1:user1" + + def test_without_user_is_unique(self): + k1 = OpenResponsesBridge._derive_session_key("", "m") + k2 = OpenResponsesBridge._derive_session_key("", "m") + assert k1 != k2 + assert k1.startswith("agent:m:") + + def test_default_model(self): + key = OpenResponsesBridge._derive_session_key("u1", "") + assert key == "agent:default:u1" + + +class TestChatBridgeResolveSessionKey: + """``ChatCompletionsBridge`` requires an explicit session key: the + ``X-AReaL-Session-Key`` header wins, else it is derived from the OpenAI + ``user`` field; with neither, the resolver returns ``None`` so the request + is rejected with ``400`` instead of being minted a random key.""" + + @staticmethod + def _request(header: str | None): + """Minimal stand-in for a Starlette ``Request`` exposing ``headers.get``.""" + + class _Headers(dict): + def get(self, key, default=None): + return super().get(key, default) + + class _Req: + def __init__(self, key): + self.headers = _Headers({SESSION_KEY_HEADER: key} if key else {}) + + return _Req(header) + + def test_explicit_header_wins(self): + req = self._request("sess-1") + key = ChatCompletionsBridge._resolve_session_key( + req, {"user": "u1", "model": "m"} + ) + assert key == "sess-1" + + def test_derived_from_user(self): + req = self._request(None) + key = ChatCompletionsBridge._resolve_session_key( + req, {"user": "u1", "model": "m"} + ) + assert key == "chat:m:u1" + + def test_user_default_model(self): + req = self._request(None) + key = ChatCompletionsBridge._resolve_session_key(req, {"user": "u1"}) + assert key == "chat:default:u1" + + def test_neither_returns_none(self): + req = self._request(None) + key = ChatCompletionsBridge._resolve_session_key(req, {"model": "m"}) + assert key is None + + +class TestBridgeBuildOutputItems: + """The bridge translates a structured turn result into OpenAI Responses + output items: a non-empty summary → assistant ``message``; each + ``tool_call`` event → a ``function_call`` item.""" + + def test_summary_becomes_message_item(self): + items = OpenResponsesBridge._build_output_items({"summary": "hello"}) + assert items == [ + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "hello"}], + } + ] + + def test_tool_call_events_become_function_calls(self): + result = { + "summary": "", + "events": [ + {"type": "tool_call", "name": "lookup", "args": '{"id": "1"}'}, + {"type": "tool_result", "name": "lookup", "result": "ok"}, + {"type": "delta", "text": "ignored"}, + ], + } + items = OpenResponsesBridge._build_output_items(result) + assert items == [ + {"type": "function_call", "name": "lookup", "arguments": '{"id": "1"}'} + ] + + def test_empty_result_yields_no_items(self): + assert OpenResponsesBridge._build_output_items({}) == [] + + +def _parse_sse(raw: bytes) -> list[dict]: + events = [] + for block in raw.decode().split("\n\n"): + block = block.strip() + if block.startswith("data: "): + events.append(json.loads(block[len("data: ") :])) + return events + + +class TestBridgeResponsesSSE: + """``/v1/responses`` streaming (A): the bridge re-encodes the collected + structured output into an OpenAI Responses-format SSE event stream + (``response.created`` → ``output_text`` deltas → ``response.completed``).""" + + @pytest.mark.asyncio + async def test_sse_reencodes_collected_output(self): + output_items = OpenResponsesBridge._build_output_items({"summary": "hi there"}) + raw = b"" + async for chunk in OpenResponsesBridge._responses_sse( + "resp-1", "GLM-5", output_items, {"k": "v"} + ): + raw += chunk + events = _parse_sse(raw) + + types = [e["type"] for e in events] + assert types[0] == "response.created" + assert "response.output_text.delta" in types + assert "response.output_text.done" in types + assert types[-1] == "response.completed" + + delta = next(e for e in events if e["type"] == "response.output_text.delta") + assert delta["delta"] == "hi there" + + completed = events[-1]["response"] + assert completed["status"] == "completed" + assert completed["model"] == "GLM-5" + assert completed["output"] == output_items + assert completed["metadata"] == {"k": "v"} + + @pytest.mark.asyncio + async def test_sse_without_message_still_completes(self): + # Tool-call-only output (no assistant message) emits no text deltas but + # still brackets the stream with created/completed. + output_items = OpenResponsesBridge._build_output_items( + {"events": [{"type": "tool_call", "name": "f", "args": "{}"}]} + ) + raw = b"" + async for chunk in OpenResponsesBridge._responses_sse( + "resp-2", "GLM-5", output_items, {} + ): + raw += chunk + events = _parse_sse(raw) + types = [e["type"] for e in events] + assert types == ["response.created", "response.completed"] + assert events[-1]["response"]["output"] == output_items + + +class TestBridgeStreamingEndToEnd: + """``handle_request`` with ``stream=True`` runs a structured collect turn, + then returns a ``text/event-stream`` re-encoding it (A), echoing the + resolved session key on the response header.""" + + @staticmethod + def _patched_post(turn_result): + original_post = httpx.AsyncClient.post + + async def patched_post(self, url, **kwargs): + url = str(url) + if url.endswith("/route"): + return httpx.Response( + 200, + json={"data_proxy_addr": "http://proxy1"}, + request=httpx.Request("POST", url), + ) + if url.endswith("/turn"): + return httpx.Response( + 200, json=turn_result, request=httpx.Request("POST", url) + ) + return await original_post(self, url, **kwargs) + + return patched_post + + @pytest.mark.asyncio + async def test_stream_response_is_sse(self): + from fastapi import FastAPI + + from areal.v2.agent_service.gateway.bridge import ( + mount_bridge, + ) + + bridge = OpenResponsesBridge(router_addr="http://router") + app = FastAPI() + mount_bridge(app, bridge) + turn_result = {"summary": "streamed reply", "events": [], "metadata": {}} + + with patch.object(httpx.AsyncClient, "post", self._patched_post(turn_result)): + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient( + transport=transport, base_url="http://gw" + ) as client: + resp = await client.post( + "/v1/responses", + headers={**_AUTH, "X-AReaL-Session-Key": "sess-1"}, + json={ + "model": "GLM-5", + "stream": True, + "input": [{"type": "message", "content": "hi"}], + }, + ) + assert resp.status_code == 200 + assert resp.headers["content-type"].startswith("text/event-stream") + assert resp.headers.get("X-AReaL-Session-Key") == "sess-1" + events = _parse_sse(resp.content) + types = [e["type"] for e in events] + assert types[0] == "response.created" + assert types[-1] == "response.completed" + delta = next(e for e in events if e["type"] == "response.output_text.delta") + assert delta["delta"] == "streamed reply" + + @pytest.mark.asyncio + async def test_non_stream_response_is_json(self): + from fastapi import FastAPI + + from areal.v2.agent_service.gateway.bridge import ( + mount_bridge, + ) + + bridge = OpenResponsesBridge(router_addr="http://router") + app = FastAPI() + mount_bridge(app, bridge) + turn_result = {"summary": "plain reply", "events": [], "metadata": {}} + + with patch.object(httpx.AsyncClient, "post", self._patched_post(turn_result)): + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient( + transport=transport, base_url="http://gw" + ) as client: + resp = await client.post( + "/v1/responses", + headers={**_AUTH, "X-AReaL-Session-Key": "sess-2"}, + json={ + "model": "GLM-5", + "input": [{"type": "message", "content": "hi"}], + }, + ) + assert resp.status_code == 200 + assert resp.headers["content-type"].startswith("application/json") + body = resp.json() + assert body["object"] == "response" + assert body["status"] == "completed" + assert body["output"][0]["content"][0]["text"] == "plain reply" diff --git a/tests/experimental/agent_service/test_protocol.py b/tests/v2/agent_service/test_protocol.py similarity index 99% rename from tests/experimental/agent_service/test_protocol.py rename to tests/v2/agent_service/test_protocol.py index b2c3ca29ce..c34df8dce6 100644 --- a/tests/experimental/agent_service/test_protocol.py +++ b/tests/v2/agent_service/test_protocol.py @@ -6,7 +6,7 @@ import pytest -from areal.experimental.agent_service.protocol import ( +from areal.v2.agent_service.protocol import ( EventFrame, FrameType, RequestFrame, diff --git a/tests/v2/cli/__init__.py b/tests/v2/cli/__init__.py new file mode 100644 index 0000000000..9881313609 --- /dev/null +++ b/tests/v2/cli/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/v2/cli/test_agent_cli_parser.py b/tests/v2/cli/test_agent_cli_parser.py new file mode 100644 index 0000000000..f2e013de63 --- /dev/null +++ b/tests/v2/cli/test_agent_cli_parser.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from click.testing import CliRunner + +from areal.v2.cli.cli import cli + + +def test_agent_help_contains_expected_commands(): + result = CliRunner().invoke(cli, ["agent", "--help"]) + + assert result.exit_code == 0 + assert "run" in result.output + assert "stop" in result.output + assert "status" in result.output + assert "ps" in result.output + assert "logs" in result.output + + +def test_agent_help_does_not_register_session_commands(): + runner = CliRunner() + + for verb in ("new_session", "switch_session", "chat", "reward"): + result = runner.invoke(cli, ["agent", verb]) + assert result.exit_code != 0 + assert "No such command" in result.output + + +def test_agent_run_does_not_accept_initial_session_key(): + result = CliRunner().invoke(cli, ["agent", "run", "--help"]) + + assert result.exit_code == 0 + assert "--session-key" not in result.output diff --git a/tests/v2/cli/test_agent_cli_run.py b/tests/v2/cli/test_agent_cli_run.py new file mode 100644 index 0000000000..b4502ac686 --- /dev/null +++ b/tests/v2/cli/test_agent_cli_run.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from areal.v2.cli.agent.commands import run +from areal.v2.cli.agent.state import ( + PairState, + ProcessState, + ServiceState, + store, +) + + +def test_run_persists_service_state(tmp_path, monkeypatch): + monkeypatch.setenv("AREAL_HOME", str(tmp_path)) + + def fake_launch_agent_stack(**kwargs): + service = kwargs["service"] + return ServiceState( + service=service, + launch_mode="detached", + agent=kwargs["agent"], + admin_api_key=kwargs["admin_api_key"], + gateway=ProcessState( + component="gateway", + pid=11, + url="http://127.0.0.1:1", + log_file="gateway.log", + ), + router=ProcessState( + component="router", + pid=12, + url="http://127.0.0.1:2", + log_file="router.log", + ), + pairs=[ + PairState( + index=0, + worker=ProcessState( + component="worker-0", + pid=13, + url="http://127.0.0.1:3", + log_file="worker.log", + ), + data_proxy=ProcessState( + component="proxy-0", + pid=14, + url="http://127.0.0.1:4", + log_file="proxy.log", + ), + ) + ], + session_timeout=kwargs["session_timeout"], + health_poll_interval=kwargs["health_poll_interval"], + drain_timeout=kwargs["drain_timeout"], + ) + + monkeypatch.setattr(run, "launch_agent_stack", fake_launch_agent_stack) + + rc = run.do_run( + agent="pkg.Agent", + service="svc", + num_pairs=1, + admin_api_key="agent-admin", + setup_timeout=1.0, + health_poll_interval=1.0, + drain_timeout=1.0, + session_timeout=60.0, + log_level="info", + force=False, + ) + + assert rc == 0 + assert store.service_state_path("svc").exists() diff --git a/tests/v2/cli/test_agent_cli_state.py b/tests/v2/cli/test_agent_cli_state.py new file mode 100644 index 0000000000..e60e5119d5 --- /dev/null +++ b/tests/v2/cli/test_agent_cli_state.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from areal.v2.cli.agent.config import load_click_default_map +from areal.v2.cli.agent.state import ( + ProcessState, + ServiceState, + store, +) + + +def test_service_state_round_trip(tmp_path, monkeypatch): + monkeypatch.setenv("AREAL_HOME", str(tmp_path)) + state = ServiceState( + service="svc", + launch_mode="detached", + agent="pkg.Agent", + admin_api_key="key", + gateway=ProcessState( + component="gateway", + pid=11, + url="http://127.0.0.1:1", + log_file="gateway.log", + ), + router=ProcessState( + component="router", + pid=12, + url="http://127.0.0.1:2", + log_file="router.log", + ), + pairs=[], + ) + + state.save() + + loaded = ServiceState.load("svc") + assert loaded.service == "svc" + assert loaded.gateway.pid == 11 + assert [pid for _, h in loaded.components() if (pid := h.pid) > 0] == [11, 12] + assert store.service_state_path("svc").exists() + + +def test_config_maps_toml_sections_to_click_defaults(tmp_path, monkeypatch): + monkeypatch.setenv("AREAL_HOME", str(tmp_path / "home")) + config_dir = tmp_path / "home" / "agent" + config_dir.mkdir(parents=True) + (config_dir / "config.toml").write_text( + "[default]\nadmin_api_key = 'base'\n[run]\nagent = 'pkg.Agent'\nnum_pairs = 4\n" + ) + extra = tmp_path / "extra.toml" + extra.write_text("[default]\nadmin_api_key = 'override'\n") + + default_map = load_click_default_map(extra=extra) + + assert default_map["run"]["admin_api_key"] == "override" + assert default_map["run"]["agent"] == "pkg.Agent" + assert default_map["run"]["num_pairs"] == 4 diff --git a/tests/v2/cli/test_inference_cli_lifecycle.py b/tests/v2/cli/test_inference_cli_lifecycle.py new file mode 100644 index 0000000000..44cd4a45f0 --- /dev/null +++ b/tests/v2/cli/test_inference_cli_lifecycle.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json + +from areal.v2.cli.inference.scheduler import TaskHandle +from areal.v2.cli.inference.state import ( + ModelEntry, + ModelReplica, + ModelState, + ServiceState, +) + + +def _service_state(gateway_pid: int = 30, router_pid: int = 40) -> ServiceState: + return ServiceState( + service="svc", + backend="local", + gateway_handle=TaskHandle( + host="127.0.0.1", ports=[8080], gpu_devices=[], ref={"pid": gateway_pid} + ), + router_handle=TaskHandle( + host="127.0.0.1", ports=[9000], gpu_devices=[], ref={"pid": router_pid} + ), + admin_api_key="admin", + started_at=1.0, + ) + + +def _replica(*, worker_pid: int, proxy_pid: int) -> ModelReplica: + return ModelReplica( + data_proxy=TaskHandle( + host="127.0.0.1", ports=[5001], gpu_devices=[], ref={"pid": proxy_pid} + ), + worker=TaskHandle( + host="127.0.0.1", ports=[5000], gpu_devices=[0], ref={"pid": worker_pid} + ), + ) + + +def test_terminate_runtime_state_kills_in_order(monkeypatch): + from areal.v2.cli.inference import common + from areal.v2.cli.inference.state import RuntimeState + + service_state = _service_state(gateway_pid=30, router_pid=40) + model_state = ModelState( + service="svc", + models={ + "m": ModelEntry( + backend="sglang:d1", + replicas=[_replica(worker_pid=10, proxy_pid=20)], + ) + }, + ) + calls = [] + + def fake_kill_pids(pids, grace_s): + calls.append((list(pids), grace_s)) + + monkeypatch.setattr(common, "kill_pids", fake_kill_pids) + + common.terminate_runtime_state( + RuntimeState(service_state=service_state, model_state=model_state), + grace_s=7.0, + ) + + # Data-flow order: data_proxy → worker → gateway → router + assert calls == [([20], 7.0), ([10], 7.0), ([30], 7.0), ([40], 7.0)] + + +def test_prepare_service_slot_force_recovers_raw_pids(tmp_path, monkeypatch): + from areal.v2.cli.inference import lifecycle as inf_lifecycle_mod + from areal.v2.cli.inference.lifecycle import inf_lifecycle + from areal.v2.cli.inference.state import store + + monkeypatch.setenv("AREAL_HOME", str(tmp_path)) + # Omit the ``backend`` field so ServiceState.load raises KeyError and + # force_replace_slot falls back to recover_pids_from_raw_state. + store.service_state_path("svc").write_text( + json.dumps( + { + "service": "svc", + "gateway_handle": { + "host": "127.0.0.1", + "ports": [8080], + "gpu_devices": [], + "ref": {"pid": 100}, + }, + "router_handle": { + "host": "127.0.0.1", + "ports": [9000], + "gpu_devices": [], + "ref": {"pid": 101}, + }, + "admin_api_key": "admin", + "started_at": 1.0, + } + ) + ) + store.models_state_path("svc").write_text( + json.dumps( + { + "service": "svc", + "models": { + "m": { + "backend": "sglang:d1", + "replicas": [ + { + "data_proxy": { + "host": "127.0.0.1", + "ports": [5001], + "gpu_devices": [], + "ref": {"pid": 300}, + }, + "worker": { + "host": "127.0.0.1", + "ports": [5000], + "gpu_devices": [0], + "ref": {"pid": 200}, + }, + } + ], + } + }, + } + ) + ) + calls = [] + + def fake_kill_pids(pids, grace_s): + calls.append((list(pids), grace_s)) + + monkeypatch.setattr(inf_lifecycle_mod, "kill_pids", fake_kill_pids) + + inf_lifecycle.force_replace_slot("svc", grace_s=5.0) + + assert calls == [([100, 101, 300, 200], 5.0)] + assert not store.service_state_path("svc").exists() + assert not store.models_state_path("svc").exists() + + +def test_foreground_cleanup_uses_latest_model_state(tmp_path, monkeypatch): + from areal.v2.cli.inference import common + from areal.v2.cli.inference.commands import run + + monkeypatch.setenv("AREAL_HOME", str(tmp_path)) + service_state = _service_state(gateway_pid=30, router_pid=40) + latest = ModelState(service="svc") + latest.models["later"] = ModelEntry( + backend="sglang:d1", + replicas=[_replica(worker_pid=10, proxy_pid=20)], + ) + latest.save() + calls = [] + + def fake_kill_pids(pids, grace_s): + calls.append((list(pids), grace_s)) + + monkeypatch.setattr(common, "kill_pids", fake_kill_pids) + + run._cleanup_runtime("svc", service_state, ModelState(service="svc"), grace_s=3.0) + + # Same data-flow order as terminate_runtime_state. + assert calls == [([20], 3.0), ([10], 3.0), ([30], 3.0), ([40], 3.0)] diff --git a/tests/v2/cli/test_inference_cli_parser.py b/tests/v2/cli/test_inference_cli_parser.py new file mode 100644 index 0000000000..896dfc8338 --- /dev/null +++ b/tests/v2/cli/test_inference_cli_parser.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from click.testing import CliRunner + +from areal.v2.cli.main import cli + + +def test_inference_commands_expose_service_flag(): + runner = CliRunner() + + for command in ( + "run", + "stop", + "status", + "register", + "deregister", + "models", + "logs", + ): + result = runner.invoke(cli, ["inf", command, "--help"]) + assert result.exit_code == 0 + assert "--service" in result.output + + +def test_inference_ps_lists_services_not_models(): + runner = CliRunner() + + result = runner.invoke(cli, ["inf", "ps", "--help"]) + + assert result.exit_code == 0 + assert "--all" in result.output + assert "--service" not in result.output + + +def test_inference_model_name_options_are_exposed(): + runner = CliRunner() + + register = runner.invoke(cli, ["inf", "register", "--help"]) + deregister = runner.invoke(cli, ["inf", "deregister", "--help"]) + + assert register.exit_code == 0 + assert "--model-name" in register.output + assert deregister.exit_code == 0 + assert "--model-name" in deregister.output + + +def test_inference_register_parses_model_name_option(monkeypatch): + from areal.v2.cli.inference.commands import register + + captured = {} + + def fake_register(model_name, opts, *, service=None): + captured["model_name"] = model_name + captured["opts"] = opts + captured["service"] = service + return 0 + + monkeypatch.setattr(register, "do_register", fake_register) + runner = CliRunner() + + result = runner.invoke( + cli, + [ + "inf", + "register", + "--service", + "svc", + "--model-name", + "m", + "--backend", + "sglang:d1", + "--model-path", + "/models/m", + ], + ) + + assert result.exit_code == 0 + assert captured["model_name"] == "m" + assert captured["service"] == "svc" + assert captured["opts"]["backend"] == "sglang:d1" + assert captured["opts"]["model_path"] == "/models/m" + + +def test_inference_deregister_parses_model_name_option(monkeypatch): + from areal.v2.cli.inference.commands import deregister + + captured = {} + + def fake_deregister(model_name, grace, force, *, service=None): + captured["model_name"] = model_name + captured["grace"] = grace + captured["force"] = force + captured["service"] = service + return 0 + + monkeypatch.setattr(deregister, "do_deregister", fake_deregister) + runner = CliRunner() + + result = runner.invoke( + cli, + [ + "inf", + "deregister", + "--service", + "svc", + "--model-name", + "m", + "--force", + ], + ) + + assert result.exit_code == 0 + assert captured == { + "model_name": "m", + "grace": 10.0, + "force": True, + "service": "svc", + } diff --git a/tests/v2/cli/test_inference_cli_state.py b/tests/v2/cli/test_inference_cli_state.py new file mode 100644 index 0000000000..03decf98cc --- /dev/null +++ b/tests/v2/cli/test_inference_cli_state.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import os + +from click.testing import CliRunner + +from areal.v2.cli.inference.scheduler import TaskHandle +from areal.v2.cli.inference.state import ( + ModelEntry, + ModelState, + ServiceState, + store, +) +from areal.v2.cli.main import cli + + +def _service_path(service: str): + return store.service_state_path(service) + + +def _current_path(): + return store.current_service_path() + + +def _resolve(name: str | None) -> str: + return store.resolve_service_name(name) + + +def _save_service(service: str, *, gateway_pid: int | None = None) -> None: + pid = os.getpid() if gateway_pid is None else gateway_pid + ServiceState( + service=service, + backend="local", + gateway_handle=TaskHandle( + host="127.0.0.1", ports=[8080], gpu_devices=[], ref={"pid": pid} + ), + router_handle=TaskHandle( + host="127.0.0.1", ports=[9000], gpu_devices=[], ref={"pid": pid} + ), + admin_api_key="admin", + started_at=1.0, + ).save() + + +def _placeholder_model() -> ModelEntry: + return ModelEntry(backend="sglang:d1", replicas=[]) + + +def test_service_and_model_state_are_per_service(tmp_path, monkeypatch): + monkeypatch.setenv("AREAL_HOME", str(tmp_path)) + _save_service("svc-a") + model_state = ModelState(service="svc-a") + model_state.models["m-a"] = _placeholder_model() + model_state.save() + + assert _service_path("svc-a").exists() + assert store.models_state_path("svc-a").exists() + assert _current_path().read_text().strip() == "svc-a" + + loaded_service = ServiceState.load("svc-a") + loaded_models = ModelState.load("svc-a") + assert loaded_service.service == "svc-a" + assert loaded_service.backend == "local" + assert list(loaded_models.models) == ["m-a"] + + +def test_resolve_service_uses_current_then_single_running(tmp_path, monkeypatch): + monkeypatch.setenv("AREAL_HOME", str(tmp_path)) + _save_service("current") + assert _resolve(None) == "current" + + _current_path().unlink() + assert _resolve(None) == "current" + + +def test_models_command_is_service_scoped(tmp_path, monkeypatch): + monkeypatch.setenv("AREAL_HOME", str(tmp_path)) + _save_service("svc-a") + _save_service("svc-b") + for service, model in (("svc-a", "m-a"), ("svc-b", "m-b")): + model_state = ModelState(service=service) + model_state.models[model] = _placeholder_model() + model_state.save() + + result = CliRunner().invoke(cli, ["inf", "models", "--service", "svc-a", "--json"]) + + assert result.exit_code == 0 + payload = json.loads(result.output) + assert [entry["name"] for entry in payload] == ["m-a"] + + +def test_ps_lists_services_and_stale_state(tmp_path, monkeypatch): + monkeypatch.setenv("AREAL_HOME", str(tmp_path)) + _save_service("live") + _save_service("stale", gateway_pid=-1) + + result = CliRunner().invoke(cli, ["inf", "ps", "--all", "--json"]) + + assert result.exit_code == 0 + payload = {entry["service"]: entry["status"] for entry in json.loads(result.output)} + assert payload == {"live": "running", "stale": "stale"} + + +def test_recover_pids_from_raw_state_walks_handles(tmp_path, monkeypatch): + monkeypatch.setenv("AREAL_HOME", str(tmp_path)) + # Hand-write the v1 schema so we exercise recover_pids_from_raw_state + # against TaskHandle.ref.pid nesting. Order of returned pids follows the + # JSON walk order: service file first, then models file; within each + # ModelReplica data_proxy is declared before worker. + _service_path("svc").write_text( + json.dumps( + { + "service": "svc", + "backend": "local", + "gateway_handle": { + "host": "127.0.0.1", + "ports": [8080], + "gpu_devices": [], + "ref": {"pid": 100}, + }, + "router_handle": { + "host": "127.0.0.1", + "ports": [9000], + "gpu_devices": [], + "ref": {"pid": 101}, + }, + "admin_api_key": "admin", + "started_at": 1.0, + } + ) + ) + store.models_state_path("svc").write_text( + json.dumps( + { + "service": "svc", + "models": { + "m": { + "backend": "sglang:d1", + "replicas": [ + { + "data_proxy": { + "host": "127.0.0.1", + "ports": [5001], + "gpu_devices": [], + "ref": {"pid": 300}, + }, + "worker": { + "host": "127.0.0.1", + "ports": [5000], + "gpu_devices": [0], + "ref": {"pid": 200}, + }, + } + ], + } + }, + } + ) + ) + + assert store.recover_pids_from_raw_state("svc") == [100, 101, 300, 200] diff --git a/tests/experimental/inference_service/__init__.py b/tests/v2/inference_service/__init__.py similarity index 100% rename from tests/experimental/inference_service/__init__.py rename to tests/v2/inference_service/__init__.py diff --git a/tests/experimental/inference_service/conftest.py b/tests/v2/inference_service/conftest.py similarity index 86% rename from tests/experimental/inference_service/conftest.py rename to tests/v2/inference_service/conftest.py index 087b263b4b..fde627a8be 100644 --- a/tests/experimental/inference_service/conftest.py +++ b/tests/v2/inference_service/conftest.py @@ -3,7 +3,7 @@ The top-level ``areal/__init__.py`` imports from modules that use PEP 695 syntax (``def func[T](...)``) which is only valid on Python 3.12+. When running tests on Python 3.10/3.11 we register lightweight stubs for the ``areal`` namespace packages -so that importing ``areal.experimental.inference_service.data_proxy.*`` never triggers the +so that importing ``areal.v2.inference_service.data_proxy.*`` never triggers the problematic top-level init. """ @@ -20,9 +20,10 @@ _STUB_PACKAGES = [ ("areal", os.path.join(_REPO_ROOT, "areal")), ("areal.experimental", os.path.join(_REPO_ROOT, "areal", "experimental")), + ("areal.v2", os.path.join(_REPO_ROOT, "areal", "v2")), ( - "areal.experimental.inference_service", - os.path.join(_REPO_ROOT, "areal", "experimental", "inference_service"), + "areal.v2.inference_service", + os.path.join(_REPO_ROOT, "areal", "v2", "inference_service"), ), ] diff --git a/tests/experimental/inference_service/integration_utils.py b/tests/v2/inference_service/integration_utils.py similarity index 100% rename from tests/experimental/inference_service/integration_utils.py rename to tests/v2/inference_service/integration_utils.py diff --git a/tests/experimental/inference_service/test_controller.py b/tests/v2/inference_service/test_controller.py similarity index 97% rename from tests/experimental/inference_service/test_controller.py rename to tests/v2/inference_service/test_controller.py index 573fe90de7..2116c8dd11 100644 --- a/tests/experimental/inference_service/test_controller.py +++ b/tests/v2/inference_service/test_controller.py @@ -9,10 +9,10 @@ import pytest from areal.api.cli_args import AgentConfig, InferenceEngineConfig -from areal.experimental.inference_service.controller.controller import ( +from areal.v2.inference_service.controller.controller import ( RolloutControllerV2, ) -from areal.experimental.inference_service.controller.workflow import ( +from areal.v2.inference_service.controller.workflow import ( InferenceServiceWorkflow, ) @@ -200,7 +200,7 @@ def test_resolve_workflow_rollout_workflow_class_raises(self): match="direct RolloutWorkflow classes are not supported", ): controller._resolve_workflow( - "areal.experimental.inference_service.controller.workflow.InferenceServiceWorkflow" + "areal.v2.inference_service.controller.workflow.InferenceServiceWorkflow" ) @@ -364,7 +364,7 @@ async def test_async_initialize_passes_callback_and_reward_timeout_to_data_proxy gpu=0, cpu=1, mem=1, - cmd="python -m areal.experimental.inference_service.guard", + cmd="python -m areal.v2.inference_service.guard", ), ), admin_api_key="test-admin-key", @@ -532,13 +532,13 @@ async def test_online_mode_waits_on_controller(self): with ( patch( - "areal.experimental.inference_service.controller.workflow.workflow_context" + "areal.v2.inference_service.controller.workflow.workflow_context" ) as mock_wf_ctx, patch( - "areal.experimental.inference_service.controller.workflow.stats_tracker" + "areal.v2.inference_service.controller.workflow.stats_tracker" ) as mock_st, patch( - "areal.experimental.inference_service.controller.workflow.deserialize_interactions" + "areal.v2.inference_service.controller.workflow.deserialize_interactions" ) as mock_deserialize, ): mock_deserialize.return_value = {"chatcmpl-1": mock_interaction} @@ -595,10 +595,10 @@ async def run(self, data, **kwargs): with ( patch( - "areal.experimental.inference_service.controller.workflow.workflow_context" + "areal.v2.inference_service.controller.workflow.workflow_context" ) as mock_wf_ctx, patch( - "areal.experimental.inference_service.controller.workflow.stats_tracker" + "areal.v2.inference_service.controller.workflow.stats_tracker" ) as mock_st, ): mock_http_session = AsyncMock() diff --git a/tests/experimental/inference_service/test_controller_integration.py b/tests/v2/inference_service/test_controller_integration.py similarity index 97% rename from tests/experimental/inference_service/test_controller_integration.py rename to tests/v2/inference_service/test_controller_integration.py index 89711cd5c2..875b374c22 100644 --- a/tests/experimental/inference_service/test_controller_integration.py +++ b/tests/v2/inference_service/test_controller_integration.py @@ -2,7 +2,7 @@ Requires GPU and a model. Marked @pytest.mark.slow to exclude from default CI. Run manually: - uv run pytest tests/experimental/inference_service/test_controller_integration.py -v -s + uv run pytest tests/v2/inference_service/test_controller_integration.py -v -s Backend filtering: pytest -m "not vllm" # skip vLLM tests (run SGLang only) @@ -29,7 +29,7 @@ import torch from PIL import Image -from tests.experimental.inference_service.integration_utils import ( +from tests.v2.inference_service.integration_utils import ( EXPR_NAME, TRIAL_NAME, check_server_health, @@ -341,7 +341,7 @@ def _make_gateway_controller_config( gpu=0, cpu=1, mem=4, - cmd="python -m areal.experimental.inference_service.guard", + cmd="python -m areal.v2.inference_service.guard", ), ), consumer_batch_size=8, @@ -367,7 +367,7 @@ def gateway_controller(sglang_server, model_path, tmp_path_factory): if not has_gpu(): pytest.skip("GPU required") - from areal.experimental.inference_service.controller.controller import ( + from areal.v2.inference_service.controller.controller import ( RolloutControllerV2, ) @@ -410,7 +410,7 @@ def gateway_controller_full_init(request, model_path, tmp_path_factory): backend = request.param from areal.api.cli_args import InferenceEngineConfig, SchedulingSpec - from areal.experimental.inference_service.controller.controller import ( + from areal.v2.inference_service.controller.controller import ( RolloutControllerV2, ) @@ -419,9 +419,7 @@ def gateway_controller_full_init(request, model_path, tmp_path_factory): model=model_path, backend=f"{backend}:d1", scheduling_spec=( - SchedulingSpec( - gpu=1, cmd="python -m areal.experimental.inference_service.guard" - ), + SchedulingSpec(gpu=1, cmd="python -m areal.v2.inference_service.guard"), ), consumer_batch_size=8, max_head_offpolicyness=1024, @@ -462,7 +460,7 @@ def gateway_controller_full_init_online(request, model_path, tmp_path_factory): InferenceEngineConfig, SchedulingSpec, ) - from areal.experimental.inference_service.controller.controller import ( + from areal.v2.inference_service.controller.controller import ( RolloutControllerV2, ) @@ -475,9 +473,7 @@ def gateway_controller_full_init_online(request, model_path, tmp_path_factory): mode="online", ), scheduling_spec=( - SchedulingSpec( - gpu=1, cmd="python -m areal.experimental.inference_service.guard" - ), + SchedulingSpec(gpu=1, cmd="python -m areal.v2.inference_service.guard"), ), consumer_batch_size=8, max_head_offpolicyness=1024, @@ -520,7 +516,7 @@ def gateway_controller_full_init_with_reward_timeout( InferenceEngineConfig, SchedulingSpec, ) - from areal.experimental.inference_service.controller.controller import ( + from areal.v2.inference_service.controller.controller import ( RolloutControllerV2, ) @@ -534,9 +530,7 @@ def gateway_controller_full_init_with_reward_timeout( set_reward_finish_timeout=3.0, ), scheduling_spec=( - SchedulingSpec( - gpu=1, cmd="python -m areal.experimental.inference_service.guard" - ), + SchedulingSpec(gpu=1, cmd="python -m areal.v2.inference_service.guard"), ), consumer_batch_size=8, max_head_offpolicyness=1024, @@ -579,7 +573,7 @@ def gateway_controller_full_init_vlm(request, vlm_model_path, tmp_path_factory): backend = request.param from areal.api.cli_args import InferenceEngineConfig, SchedulingSpec - from areal.experimental.inference_service.controller.controller import ( + from areal.v2.inference_service.controller.controller import ( RolloutControllerV2, ) @@ -588,9 +582,7 @@ def gateway_controller_full_init_vlm(request, vlm_model_path, tmp_path_factory): model=vlm_model_path, backend=f"{backend}:d1", scheduling_spec=( - SchedulingSpec( - gpu=1, cmd="python -m areal.experimental.inference_service.guard" - ), + SchedulingSpec(gpu=1, cmd="python -m areal.v2.inference_service.guard"), ), consumer_batch_size=8, max_head_offpolicyness=1024, diff --git a/tests/experimental/inference_service/test_controller_version.py b/tests/v2/inference_service/test_controller_version.py similarity index 98% rename from tests/experimental/inference_service/test_controller_version.py rename to tests/v2/inference_service/test_controller_version.py index bba3fa7841..af7ae18b73 100644 --- a/tests/experimental/inference_service/test_controller_version.py +++ b/tests/v2/inference_service/test_controller_version.py @@ -9,7 +9,7 @@ from unittest.mock import AsyncMock, MagicMock, patch from areal.api.cli_args import InferenceEngineConfig, SchedulingSpec -from areal.experimental.inference_service.controller.controller import ( +from areal.v2.inference_service.controller.controller import ( RolloutControllerV2, ) diff --git a/tests/experimental/inference_service/test_data_proxy_chat.py b/tests/v2/inference_service/test_data_proxy_chat.py similarity index 98% rename from tests/experimental/inference_service/test_data_proxy_chat.py rename to tests/v2/inference_service/test_data_proxy_chat.py index 2965e70270..8c241374fb 100644 --- a/tests/experimental/inference_service/test_data_proxy_chat.py +++ b/tests/v2/inference_service/test_data_proxy_chat.py @@ -9,12 +9,12 @@ import pytest import pytest_asyncio -from areal.experimental.inference_service.data_proxy.app import ( +from areal.v2.inference_service.data_proxy.app import ( _flush_ready_trajectories, create_app, ) -from areal.experimental.inference_service.data_proxy.config import DataProxyConfig -from areal.experimental.inference_service.data_proxy.session import ( +from areal.v2.inference_service.data_proxy.config import DataProxyConfig +from areal.v2.inference_service.data_proxy.session import ( SessionData, SessionStore, ) @@ -121,9 +121,9 @@ async def _mock_create(*, areal_cache=None, **kwargs): @pytest_asyncio.fixture async def client(config, mock_tokenizer, mock_areal_client): """Create app with mocked deps and yield an httpx async client.""" - from areal.experimental.inference_service.data_proxy.pause import PauseState - from areal.experimental.inference_service.inf_bridge import InfBridge - from areal.experimental.inference_service.sglang.bridge import SGLangBridgeBackend + from areal.v2.inference_service.data_proxy.pause import PauseState + from areal.v2.inference_service.inf_bridge import InfBridge + from areal.v2.inference_service.sglang.bridge import SGLangBridgeBackend app = create_app(config) # Bypass lifespan — inject mocks directly into app.state @@ -608,7 +608,7 @@ async def _mock_post_callback( return True monkeypatch.setattr( - "areal.experimental.inference_service.data_proxy.app._post_online_ready_callback", + "areal.v2.inference_service.data_proxy.app._post_online_ready_callback", _mock_post_callback, ) diff --git a/tests/experimental/inference_service/test_data_proxy_integration.py b/tests/v2/inference_service/test_data_proxy_integration.py similarity index 97% rename from tests/experimental/inference_service/test_data_proxy_integration.py rename to tests/v2/inference_service/test_data_proxy_integration.py index 8ce11f900c..0d68fc98a2 100644 --- a/tests/experimental/inference_service/test_data_proxy_integration.py +++ b/tests/v2/inference_service/test_data_proxy_integration.py @@ -1,7 +1,7 @@ """Integration tests for data proxy with a real SGLang server. Requires GPU and a model. Run with: - uv run pytest tests/experimental/inference_service/test_data_proxy_integration.py -v -s + uv run pytest tests/v2/inference_service/test_data_proxy_integration.py -v -s The test launches an SGLang server subprocess, starts the data proxy FastAPI app, and exercises the /chat/completions endpoint (full session lifecycle) and @@ -103,7 +103,7 @@ def vllm_server(): if not _has_gpu(): pytest.skip("GPU required for vLLM server") - from tests.experimental.inference_service.integration_utils import ( + from tests.v2.inference_service.integration_utils import ( launch_vllm_server, ) @@ -126,16 +126,16 @@ def model_path() -> str: def _create_data_proxy_app_with_sessions(sglang_server, model_path): """Create a fully-wired data proxy app with session support.""" - from areal.experimental.inference_service.data_proxy.app import create_app - from areal.experimental.inference_service.data_proxy.config import DataProxyConfig - from areal.experimental.inference_service.data_proxy.pause import PauseState - from areal.experimental.inference_service.data_proxy.session import SessionStore - from areal.experimental.inference_service.data_proxy.tokenizer_proxy import ( + from areal.experimental.openai.client import ArealOpenAI + from areal.v2.inference_service.data_proxy.app import create_app + from areal.v2.inference_service.data_proxy.config import DataProxyConfig + from areal.v2.inference_service.data_proxy.pause import PauseState + from areal.v2.inference_service.data_proxy.session import SessionStore + from areal.v2.inference_service.data_proxy.tokenizer_proxy import ( TokenizerProxy, ) - from areal.experimental.inference_service.inf_bridge import InfBridge - from areal.experimental.inference_service.sglang.bridge import SGLangBridgeBackend - from areal.experimental.openai.client import ArealOpenAI + from areal.v2.inference_service.inf_bridge import InfBridge + from areal.v2.inference_service.sglang.bridge import SGLangBridgeBackend config = DataProxyConfig( host="127.0.0.1", @@ -171,16 +171,16 @@ def _create_data_proxy_app_with_sessions(sglang_server, model_path): def _create_data_proxy_app_vllm(vllm_server, model_path): """Create a data proxy app backed by vLLM with session support.""" - from areal.experimental.inference_service.data_proxy.app import create_app - from areal.experimental.inference_service.data_proxy.config import DataProxyConfig - from areal.experimental.inference_service.data_proxy.pause import PauseState - from areal.experimental.inference_service.data_proxy.session import SessionStore - from areal.experimental.inference_service.data_proxy.tokenizer_proxy import ( + from areal.experimental.openai.client import ArealOpenAI + from areal.v2.inference_service.data_proxy.app import create_app + from areal.v2.inference_service.data_proxy.config import DataProxyConfig + from areal.v2.inference_service.data_proxy.pause import PauseState + from areal.v2.inference_service.data_proxy.session import SessionStore + from areal.v2.inference_service.data_proxy.tokenizer_proxy import ( TokenizerProxy, ) - from areal.experimental.inference_service.inf_bridge import InfBridge - from areal.experimental.inference_service.vllm.bridge import VLLMBridgeBackend - from areal.experimental.openai.client import ArealOpenAI + from areal.v2.inference_service.inf_bridge import InfBridge + from areal.v2.inference_service.vllm.bridge import VLLMBridgeBackend config = DataProxyConfig( host="127.0.0.1", diff --git a/tests/experimental/inference_service/test_data_proxy_pause.py b/tests/v2/inference_service/test_data_proxy_pause.py similarity index 94% rename from tests/experimental/inference_service/test_data_proxy_pause.py rename to tests/v2/inference_service/test_data_proxy_pause.py index 330702032e..32fea7781d 100644 --- a/tests/experimental/inference_service/test_data_proxy_pause.py +++ b/tests/v2/inference_service/test_data_proxy_pause.py @@ -8,10 +8,10 @@ import pytest import pytest_asyncio -from areal.experimental.inference_service.data_proxy.app import create_app -from areal.experimental.inference_service.data_proxy.config import DataProxyConfig -from areal.experimental.inference_service.data_proxy.pause import PauseState -from areal.experimental.inference_service.data_proxy.session import SessionStore +from areal.v2.inference_service.data_proxy.app import create_app +from areal.v2.inference_service.data_proxy.config import DataProxyConfig +from areal.v2.inference_service.data_proxy.pause import PauseState +from areal.v2.inference_service.data_proxy.session import SessionStore # ============================================================================= # Fixtures @@ -99,8 +99,8 @@ async def _mock_create(*, areal_cache=None, **kwargs): @pytest_asyncio.fixture async def app_client(config, mock_tokenizer, mock_areal_client): """Create an ASGI test client with all app.state attributes injected.""" - from areal.experimental.inference_service.inf_bridge import InfBridge - from areal.experimental.inference_service.sglang.bridge import SGLangBridgeBackend + from areal.v2.inference_service.inf_bridge import InfBridge + from areal.v2.inference_service.sglang.bridge import SGLangBridgeBackend app = create_app(config) diff --git a/tests/experimental/inference_service/test_data_proxy_rtensor.py b/tests/v2/inference_service/test_data_proxy_rtensor.py similarity index 97% rename from tests/experimental/inference_service/test_data_proxy_rtensor.py rename to tests/v2/inference_service/test_data_proxy_rtensor.py index 306cc769bd..2dd5923043 100644 --- a/tests/experimental/inference_service/test_data_proxy_rtensor.py +++ b/tests/v2/inference_service/test_data_proxy_rtensor.py @@ -18,11 +18,11 @@ import pytest_asyncio import torch -from areal.experimental.inference_service.data_proxy.app import create_app -from areal.experimental.inference_service.data_proxy.config import DataProxyConfig -from areal.experimental.inference_service.data_proxy.session import SessionStore from areal.infra.rpc import rtensor as rtensor_storage from areal.infra.rpc.serialization import deserialize_value, serialize_value +from areal.v2.inference_service.data_proxy.app import create_app +from areal.v2.inference_service.data_proxy.config import DataProxyConfig +from areal.v2.inference_service.data_proxy.session import SessionStore # ============================================================================= # Fixtures @@ -51,7 +51,7 @@ def config(): @pytest_asyncio.fixture async def client(config): - from areal.experimental.inference_service.data_proxy.pause import PauseState + from areal.v2.inference_service.data_proxy.pause import PauseState app = create_app(config) pause_state = PauseState() diff --git a/tests/experimental/inference_service/test_data_proxy_standalone.py b/tests/v2/inference_service/test_data_proxy_standalone.py similarity index 91% rename from tests/experimental/inference_service/test_data_proxy_standalone.py rename to tests/v2/inference_service/test_data_proxy_standalone.py index 0a29516dee..3dd723b4f1 100644 --- a/tests/experimental/inference_service/test_data_proxy_standalone.py +++ b/tests/v2/inference_service/test_data_proxy_standalone.py @@ -13,9 +13,9 @@ import pytest import pytest_asyncio -from areal.experimental.inference_service.data_proxy.app import create_app -from areal.experimental.inference_service.data_proxy.config import DataProxyConfig -from areal.experimental.inference_service.data_proxy.session import SessionStore +from areal.v2.inference_service.data_proxy.app import create_app +from areal.v2.inference_service.data_proxy.config import DataProxyConfig +from areal.v2.inference_service.data_proxy.session import SessionStore # ============================================================================= # Fixtures @@ -85,9 +85,9 @@ async def _mock_create(*, areal_cache=None, **kwargs): @pytest_asyncio.fixture async def client(config, mock_tokenizer, mock_areal_client): """Create app with mocked deps and yield an httpx async client (no auth header).""" - from areal.experimental.inference_service.data_proxy.pause import PauseState - from areal.experimental.inference_service.inf_bridge import InfBridge - from areal.experimental.inference_service.sglang.bridge import SGLangBridgeBackend + from areal.v2.inference_service.data_proxy.pause import PauseState + from areal.v2.inference_service.inf_bridge import InfBridge + from areal.v2.inference_service.sglang.bridge import SGLangBridgeBackend app = create_app(config) pause_state = PauseState() @@ -123,11 +123,11 @@ def admin_headers(): class TestStandaloneChat: def test_config_can_select_vllm_backend(self): """backend_type=vllm creates a vLLM-backed InfBridge.""" - from areal.experimental.inference_service.data_proxy.app import ( + from areal.v2.inference_service.data_proxy.app import ( _create_inf_bridge, ) - from areal.experimental.inference_service.data_proxy.pause import PauseState - from areal.experimental.inference_service.vllm.bridge import VLLMBridgeBackend + from areal.v2.inference_service.data_proxy.pause import PauseState + from areal.v2.inference_service.vllm.bridge import VLLMBridgeBackend config = DataProxyConfig( host="127.0.0.1", diff --git a/tests/experimental/inference_service/test_data_proxy_version.py b/tests/v2/inference_service/test_data_proxy_version.py similarity index 93% rename from tests/experimental/inference_service/test_data_proxy_version.py rename to tests/v2/inference_service/test_data_proxy_version.py index 6a65d1cc01..c5031b33a9 100644 --- a/tests/experimental/inference_service/test_data_proxy_version.py +++ b/tests/v2/inference_service/test_data_proxy_version.py @@ -8,10 +8,10 @@ import pytest import pytest_asyncio -from areal.experimental.inference_service.data_proxy.app import create_app -from areal.experimental.inference_service.data_proxy.config import DataProxyConfig -from areal.experimental.inference_service.data_proxy.pause import PauseState -from areal.experimental.inference_service.data_proxy.session import SessionStore +from areal.v2.inference_service.data_proxy.app import create_app +from areal.v2.inference_service.data_proxy.config import DataProxyConfig +from areal.v2.inference_service.data_proxy.pause import PauseState +from areal.v2.inference_service.data_proxy.session import SessionStore # ============================================================================= # Fixtures @@ -99,8 +99,8 @@ async def _mock_create(*, areal_cache=None, **kwargs): @pytest_asyncio.fixture async def app_client(config, mock_tokenizer, mock_areal_client): """Create an ASGI test client with all app.state attributes injected.""" - from areal.experimental.inference_service.inf_bridge import InfBridge - from areal.experimental.inference_service.sglang.bridge import SGLangBridgeBackend + from areal.v2.inference_service.inf_bridge import InfBridge + from areal.v2.inference_service.sglang.bridge import SGLangBridgeBackend app = create_app(config) diff --git a/tests/experimental/inference_service/test_examples.py b/tests/v2/inference_service/test_examples.py similarity index 100% rename from tests/experimental/inference_service/test_examples.py rename to tests/v2/inference_service/test_examples.py diff --git a/tests/experimental/inference_service/test_external_model.py b/tests/v2/inference_service/test_external_model.py similarity index 96% rename from tests/experimental/inference_service/test_external_model.py rename to tests/v2/inference_service/test_external_model.py index 8a6430f592..93dc8bd580 100644 --- a/tests/experimental/inference_service/test_external_model.py +++ b/tests/v2/inference_service/test_external_model.py @@ -8,28 +8,28 @@ import pytest import pytest_asyncio -from areal.experimental.inference_service.data_proxy.app import ( +from areal.v2.inference_service.data_proxy.app import ( create_app as create_data_proxy_app, ) -from areal.experimental.inference_service.data_proxy.config import DataProxyConfig -from areal.experimental.inference_service.data_proxy.session import SessionStore -from areal.experimental.inference_service.gateway.app import ( +from areal.v2.inference_service.data_proxy.config import DataProxyConfig +from areal.v2.inference_service.data_proxy.session import SessionStore +from areal.v2.inference_service.gateway.app import ( create_app as create_gateway_app, ) -from areal.experimental.inference_service.gateway.config import GatewayConfig -from areal.experimental.inference_service.gateway.streaming import ( +from areal.v2.inference_service.gateway.config import GatewayConfig +from areal.v2.inference_service.gateway.streaming import ( RouterKeyRejectedError, ) -from areal.experimental.inference_service.router.app import ( +from areal.v2.inference_service.router.app import ( create_app as create_router_app, ) -from areal.experimental.inference_service.router.config import RouterConfig -from areal.experimental.inference_service.router.state import ModelRegistry +from areal.v2.inference_service.router.config import RouterConfig +from areal.v2.inference_service.router.state import ModelRegistry ADMIN_KEY = "test-admin-key" SESSION_KEY = "session-key-abc123" WORKER_ADDR = "http://worker-1:18082" -ROUTER_MODULE = "areal.experimental.inference_service.gateway.app" +ROUTER_MODULE = "areal.v2.inference_service.gateway.app" def admin_headers() -> dict[str, str]: @@ -411,9 +411,9 @@ def mock_areal_client(): @pytest_asyncio.fixture async def data_proxy_app(data_proxy_config, mock_tokenizer, mock_areal_client): - from areal.experimental.inference_service.data_proxy.pause import PauseState - from areal.experimental.inference_service.inf_bridge import InfBridge - from areal.experimental.inference_service.sglang.bridge import SGLangBridgeBackend + from areal.v2.inference_service.data_proxy.pause import PauseState + from areal.v2.inference_service.inf_bridge import InfBridge + from areal.v2.inference_service.sglang.bridge import SGLangBridgeBackend app = create_data_proxy_app(data_proxy_config) pause_state = PauseState() @@ -555,7 +555,7 @@ def stream(self, method, url, json=None, headers=None): return _FakeStreamResponse() monkeypatch.setattr( - "areal.experimental.inference_service.data_proxy.app.httpx.AsyncClient", + "areal.v2.inference_service.data_proxy.app.httpx.AsyncClient", _FakeClient, ) diff --git a/tests/experimental/inference_service/test_external_model_integration.py b/tests/v2/inference_service/test_external_model_integration.py similarity index 94% rename from tests/experimental/inference_service/test_external_model_integration.py rename to tests/v2/inference_service/test_external_model_integration.py index 9d2b7ad048..b8a13eb2cc 100644 --- a/tests/experimental/inference_service/test_external_model_integration.py +++ b/tests/v2/inference_service/test_external_model_integration.py @@ -9,26 +9,26 @@ from fastapi import FastAPI, Request from fastapi.responses import JSONResponse -from areal.experimental.inference_service.data_proxy.app import ( +from areal.v2.inference_service.data_proxy.app import ( create_app as create_data_proxy_app, ) -from areal.experimental.inference_service.data_proxy.config import DataProxyConfig -from areal.experimental.inference_service.data_proxy.session import SessionStore -from areal.experimental.inference_service.gateway.app import ( +from areal.v2.inference_service.data_proxy.config import DataProxyConfig +from areal.v2.inference_service.data_proxy.session import SessionStore +from areal.v2.inference_service.gateway.app import ( create_app as create_gateway_app, ) -from areal.experimental.inference_service.gateway.config import GatewayConfig -from areal.experimental.inference_service.gateway.streaming import ( +from areal.v2.inference_service.gateway.config import GatewayConfig +from areal.v2.inference_service.gateway.streaming import ( RouterKeyRejectedError, ) -from areal.experimental.inference_service.router.app import ( +from areal.v2.inference_service.router.app import ( create_app as create_router_app, ) -from areal.experimental.inference_service.router.config import RouterConfig +from areal.v2.inference_service.router.config import RouterConfig ADMIN_KEY = "test-admin-key" WORKER_ADDR = "http://worker-1:18082" -ROUTER_MODULE = "areal.experimental.inference_service.gateway.app" +ROUTER_MODULE = "areal.v2.inference_service.gateway.app" def admin_headers() -> dict[str, str]: @@ -339,7 +339,7 @@ async def post(self, url, json=None, headers=None): new=AsyncMock(side_effect=_forward_request), ), patch( - "areal.experimental.inference_service.data_proxy.app.httpx.AsyncClient", + "areal.v2.inference_service.data_proxy.app.httpx.AsyncClient", _ExternalClient, ), ): diff --git a/tests/experimental/inference_service/test_gateway.py b/tests/v2/inference_service/test_gateway.py similarity index 98% rename from tests/experimental/inference_service/test_gateway.py rename to tests/v2/inference_service/test_gateway.py index 5f9830a48b..a57f209914 100644 --- a/tests/experimental/inference_service/test_gateway.py +++ b/tests/v2/inference_service/test_gateway.py @@ -12,9 +12,9 @@ import pytest import pytest_asyncio -from areal.experimental.inference_service.gateway.app import create_app -from areal.experimental.inference_service.gateway.config import GatewayConfig -from areal.experimental.inference_service.gateway.streaming import ( +from areal.v2.inference_service.gateway.app import create_app +from areal.v2.inference_service.gateway.config import GatewayConfig +from areal.v2.inference_service.gateway.streaming import ( RouterKeyRejectedError, RouterUnreachableError, ) @@ -29,7 +29,7 @@ WORKER_ADDR_2 = "http://worker-2:18082" WORKER_ADDR_3 = "http://worker-3:18082" -MODULE = "areal.experimental.inference_service.gateway.app" +MODULE = "areal.v2.inference_service.gateway.app" @pytest.fixture diff --git a/tests/experimental/inference_service/test_gateway_integration.py b/tests/v2/inference_service/test_gateway_integration.py similarity index 97% rename from tests/experimental/inference_service/test_gateway_integration.py rename to tests/v2/inference_service/test_gateway_integration.py index 52be4b1444..cad2b28a24 100644 --- a/tests/experimental/inference_service/test_gateway_integration.py +++ b/tests/v2/inference_service/test_gateway_integration.py @@ -1,7 +1,7 @@ """Full-stack integration test: client → Gateway → Router → Data Proxy → SGLang. Requires GPU and a model. Run with: - uv run pytest tests/experimental/inference_service/test_gateway_integration.py -v -s + uv run pytest tests/v2/inference_service/test_gateway_integration.py -v -s The test launches: 1. A real SGLang server (GPU subprocess) @@ -136,7 +136,7 @@ def vllm_server(): if not _has_gpu(): pytest.skip("GPU required for vLLM server") - from tests.experimental.inference_service.integration_utils import ( + from tests.v2.inference_service.integration_utils import ( launch_vllm_server, ) @@ -169,10 +169,10 @@ def gateway_stack(sglang_server, model_path): gateway_addr = f"http://{bind_host}:{gateway_port}" # --- Create Data Proxy app --- - from areal.experimental.inference_service.data_proxy.app import ( + from areal.v2.inference_service.data_proxy.app import ( create_app as create_dp_app, ) - from areal.experimental.inference_service.data_proxy.config import DataProxyConfig + from areal.v2.inference_service.data_proxy.config import DataProxyConfig dp_config = DataProxyConfig( host=bind_host, @@ -185,10 +185,10 @@ def gateway_stack(sglang_server, model_path): dp_app = create_dp_app(dp_config) # --- Create Router app --- - from areal.experimental.inference_service.router.app import ( + from areal.v2.inference_service.router.app import ( create_app as create_router_app, ) - from areal.experimental.inference_service.router.config import RouterConfig + from areal.v2.inference_service.router.config import RouterConfig router_config = RouterConfig( host=bind_host, @@ -201,10 +201,10 @@ def gateway_stack(sglang_server, model_path): router_app = create_router_app(router_config) # --- Create Gateway app --- - from areal.experimental.inference_service.gateway.app import ( + from areal.v2.inference_service.gateway.app import ( create_app as create_gw_app, ) - from areal.experimental.inference_service.gateway.config import GatewayConfig + from areal.v2.inference_service.gateway.config import GatewayConfig gw_config = GatewayConfig( host=bind_host, @@ -798,18 +798,18 @@ async def pause_then_resume(): @pytest.fixture(scope="module") def gateway_stack_vllm(vllm_server, model_path): """Launch the full gateway stack backed by vLLM on free ports.""" - from areal.experimental.inference_service.data_proxy.app import ( + from areal.v2.inference_service.data_proxy.app import ( create_app as create_dp_app, ) - from areal.experimental.inference_service.data_proxy.config import DataProxyConfig - from areal.experimental.inference_service.gateway.app import ( + from areal.v2.inference_service.data_proxy.config import DataProxyConfig + from areal.v2.inference_service.gateway.app import ( create_app as create_gw_app, ) - from areal.experimental.inference_service.gateway.config import GatewayConfig - from areal.experimental.inference_service.router.app import ( + from areal.v2.inference_service.gateway.config import GatewayConfig + from areal.v2.inference_service.router.app import ( create_app as create_router_app, ) - from areal.experimental.inference_service.router.config import RouterConfig + from areal.v2.inference_service.router.config import RouterConfig bind_host = "127.0.0.1" data_proxy_port = _find_free_port() diff --git a/tests/experimental/inference_service/test_guard.py b/tests/v2/inference_service/test_guard.py similarity index 97% rename from tests/experimental/inference_service/test_guard.py rename to tests/v2/inference_service/test_guard.py index 8922ba9e54..805a8340ef 100644 --- a/tests/experimental/inference_service/test_guard.py +++ b/tests/v2/inference_service/test_guard.py @@ -1,4 +1,4 @@ -"""Unit tests for RPCGuard Flask app (areal.experimental.inference_service.guard.app). +"""Unit tests for RPCGuard Flask app (areal.v2.inference_service.guard.app). Tests all 4 endpoints (/health, /alloc_ports, /fork, /kill_forked_worker) and the cleanup_forked_children() function using Flask test client with @@ -12,8 +12,8 @@ import pytest -from areal.experimental.inference_service.guard import app as guard_module -from areal.experimental.inference_service.guard.app import app, cleanup_forked_children +from areal.v2.inference_service.guard import app as guard_module +from areal.v2.inference_service.guard.app import app, cleanup_forked_children GUARD_APP = "areal.infra.rpc.guard.app" @@ -133,7 +133,7 @@ def test_fork_raw_cmd_passes_command_as_is(self, mock_run, client): raw = [ "python", "-m", - "areal.experimental.inference_service.sglang.launch_server", + "areal.v2.inference_service.sglang.launch_server", "--model", "test-model", ] diff --git a/tests/experimental/inference_service/test_image_input.py b/tests/v2/inference_service/test_image_input.py similarity index 98% rename from tests/experimental/inference_service/test_image_input.py rename to tests/v2/inference_service/test_image_input.py index 1088c7b6e7..ce0821fc0a 100644 --- a/tests/experimental/inference_service/test_image_input.py +++ b/tests/v2/inference_service/test_image_input.py @@ -19,14 +19,14 @@ from areal.api.cli_args import GenerationHyperparameters from areal.api.io_struct import ModelRequest -from areal.experimental.inference_service.data_proxy.pause import PauseState -from areal.experimental.inference_service.inf_bridge import InfBridge -from areal.experimental.inference_service.sglang.bridge import SGLangBridgeBackend -from areal.experimental.inference_service.vllm.bridge import VLLMBridgeBackend from areal.experimental.openai.client import ( _build_messages_list, _extract_images_from_messages, ) +from areal.v2.inference_service.data_proxy.pause import PauseState +from areal.v2.inference_service.inf_bridge import InfBridge +from areal.v2.inference_service.sglang.bridge import SGLangBridgeBackend +from areal.v2.inference_service.vllm.bridge import VLLMBridgeBackend # --------------------------------------------------------------------------- # Real image fixtures @@ -663,9 +663,9 @@ async def _mock_create(*, areal_cache=None, **kwargs): @pytest_asyncio.fixture async def image_test_client(): """Data proxy test client wired with an image-capturing mock.""" - from areal.experimental.inference_service.data_proxy.app import create_app - from areal.experimental.inference_service.data_proxy.config import DataProxyConfig - from areal.experimental.inference_service.data_proxy.session import SessionStore + from areal.v2.inference_service.data_proxy.app import create_app + from areal.v2.inference_service.data_proxy.config import DataProxyConfig + from areal.v2.inference_service.data_proxy.session import SessionStore config = DataProxyConfig( host="127.0.0.1", diff --git a/tests/experimental/inference_service/test_inf_bridge.py b/tests/v2/inference_service/test_inf_bridge.py similarity index 98% rename from tests/experimental/inference_service/test_inf_bridge.py rename to tests/v2/inference_service/test_inf_bridge.py index 2216b92d1e..aad70d546b 100644 --- a/tests/experimental/inference_service/test_inf_bridge.py +++ b/tests/v2/inference_service/test_inf_bridge.py @@ -10,10 +10,10 @@ from areal.api.cli_args import GenerationHyperparameters from areal.api.io_struct import ModelRequest, ModelResponse, get_versioned_lora_name -from areal.experimental.inference_service.data_proxy.pause import PauseState -from areal.experimental.inference_service.inf_bridge import InfBridge -from areal.experimental.inference_service.sglang.bridge import SGLangBridgeBackend -from areal.experimental.inference_service.vllm.bridge import VLLMBridgeBackend +from areal.v2.inference_service.data_proxy.pause import PauseState +from areal.v2.inference_service.inf_bridge import InfBridge +from areal.v2.inference_service.sglang.bridge import SGLangBridgeBackend +from areal.v2.inference_service.vllm.bridge import VLLMBridgeBackend # ============================================================================= # Helpers diff --git a/tests/experimental/inference_service/test_ipv6_entrypoints.py b/tests/v2/inference_service/test_ipv6_entrypoints.py similarity index 94% rename from tests/experimental/inference_service/test_ipv6_entrypoints.py rename to tests/v2/inference_service/test_ipv6_entrypoints.py index d382042c06..2b7d9b37dd 100644 --- a/tests/experimental/inference_service/test_ipv6_entrypoints.py +++ b/tests/v2/inference_service/test_ipv6_entrypoints.py @@ -5,7 +5,7 @@ def test_data_proxy_main_formats_ipv6_serving_addr(): - from areal.experimental.inference_service.data_proxy import ( + from areal.v2.inference_service.data_proxy import ( __main__ as data_proxy_main, ) @@ -45,7 +45,7 @@ def test_data_proxy_main_formats_ipv6_serving_addr(): def test_guard_main_registers_ipv6_worker_addr(): - from areal.experimental.inference_service.guard import __main__ as guard_main + from areal.v2.inference_service.guard import __main__ as guard_main args = argparse.Namespace( port=0, diff --git a/tests/experimental/inference_service/test_online_stack.py b/tests/v2/inference_service/test_online_stack.py similarity index 92% rename from tests/experimental/inference_service/test_online_stack.py rename to tests/v2/inference_service/test_online_stack.py index c8350c8efd..d39b0c1a5f 100644 --- a/tests/experimental/inference_service/test_online_stack.py +++ b/tests/v2/inference_service/test_online_stack.py @@ -13,21 +13,21 @@ import pytest import pytest_asyncio -from areal.experimental.inference_service.data_proxy.app import ( +from areal.v2.inference_service.data_proxy.app import ( create_app as create_dp_app, ) -from areal.experimental.inference_service.data_proxy.config import DataProxyConfig -from areal.experimental.inference_service.data_proxy.session import SessionStore -from areal.experimental.inference_service.gateway.app import create_app as create_gw_app -from areal.experimental.inference_service.gateway.config import GatewayConfig -from areal.experimental.inference_service.gateway.streaming import ( +from areal.v2.inference_service.data_proxy.config import DataProxyConfig +from areal.v2.inference_service.data_proxy.session import SessionStore +from areal.v2.inference_service.gateway.app import create_app as create_gw_app +from areal.v2.inference_service.gateway.config import GatewayConfig +from areal.v2.inference_service.gateway.streaming import ( RouterKeyRejectedError, RouterUnreachableError, ) -from areal.experimental.inference_service.router.app import ( +from areal.v2.inference_service.router.app import ( create_app as create_router_app, ) -from areal.experimental.inference_service.router.config import RouterConfig +from areal.v2.inference_service.router.config import RouterConfig ADMIN_KEY = "areal-admin-key" DATA_PROXY_ADDR = "http://data-proxy" @@ -88,10 +88,10 @@ async def _mock_create(*, areal_cache=None, **kwargs): @pytest_asyncio.fixture async def online_stack(monkeypatch): - import areal.experimental.inference_service.gateway.app as gateway_app_module - from areal.experimental.inference_service.data_proxy.pause import PauseState - from areal.experimental.inference_service.inf_bridge import InfBridge - from areal.experimental.inference_service.sglang.bridge import SGLangBridgeBackend + import areal.v2.inference_service.gateway.app as gateway_app_module + from areal.v2.inference_service.data_proxy.pause import PauseState + from areal.v2.inference_service.inf_bridge import InfBridge + from areal.v2.inference_service.sglang.bridge import SGLangBridgeBackend dp_config = DataProxyConfig( host="127.0.0.1", diff --git a/tests/experimental/inference_service/test_router.py b/tests/v2/inference_service/test_router.py similarity index 98% rename from tests/experimental/inference_service/test_router.py rename to tests/v2/inference_service/test_router.py index e03c35f442..0199b842eb 100644 --- a/tests/experimental/inference_service/test_router.py +++ b/tests/v2/inference_service/test_router.py @@ -10,14 +10,14 @@ import pytest import pytest_asyncio -from areal.experimental.inference_service.router.app import create_app -from areal.experimental.inference_service.router.config import RouterConfig -from areal.experimental.inference_service.router.state import ( +from areal.v2.inference_service.router.app import create_app +from areal.v2.inference_service.router.config import RouterConfig +from areal.v2.inference_service.router.state import ( SessionRegistry, WorkerInfo, WorkerRegistry, ) -from areal.experimental.inference_service.router.strategies import ( +from areal.v2.inference_service.router.strategies import ( RoundRobinStrategy, get_strategy, ) diff --git a/tests/experimental/inference_service/test_sglang_multinode.py b/tests/v2/inference_service/test_sglang_multinode.py similarity index 100% rename from tests/experimental/inference_service/test_sglang_multinode.py rename to tests/v2/inference_service/test_sglang_multinode.py diff --git a/tests/experimental/inference_service/test_vllm_multinode.py b/tests/v2/inference_service/test_vllm_multinode.py similarity index 100% rename from tests/experimental/inference_service/test_vllm_multinode.py rename to tests/v2/inference_service/test_vllm_multinode.py diff --git a/tests/experimental/training_service/__init__.py b/tests/v2/training_service/__init__.py similarity index 100% rename from tests/experimental/training_service/__init__.py rename to tests/v2/training_service/__init__.py diff --git a/tests/experimental/training_service/fake_train_engine.py b/tests/v2/training_service/fake_train_engine.py similarity index 100% rename from tests/experimental/training_service/fake_train_engine.py rename to tests/v2/training_service/fake_train_engine.py diff --git a/tests/experimental/training_service/test_controller_integration.py b/tests/v2/training_service/test_controller_integration.py similarity index 99% rename from tests/experimental/training_service/test_controller_integration.py rename to tests/v2/training_service/test_controller_integration.py index 33b5286778..a6edc605b7 100644 --- a/tests/experimental/training_service/test_controller_integration.py +++ b/tests/v2/training_service/test_controller_integration.py @@ -16,11 +16,11 @@ TrainEngineConfig, ) from areal.api.io_struct import FinetuneSpec, SaveLoadMeta -from areal.experimental.training_service.controller.controller import ( - GatewayTrainController, -) from areal.infra.platforms import current_platform from areal.infra.scheduler.local import LocalScheduler +from areal.v2.training_service.controller.controller import ( + GatewayTrainController, +) LOCAL_MODEL_PATH = "/storage/openpsi/models/Qwen__Qwen3-0.6B/" LOCAL_MOE_MODEL_PATH = "/storage/openpsi/models/Qwen__Qwen3-30B-A3B/" diff --git a/tests/experimental/training_service/test_controller_unit.py b/tests/v2/training_service/test_controller_unit.py similarity index 97% rename from tests/experimental/training_service/test_controller_unit.py rename to tests/v2/training_service/test_controller_unit.py index 4495fd942d..a489e0c613 100644 --- a/tests/experimental/training_service/test_controller_unit.py +++ b/tests/v2/training_service/test_controller_unit.py @@ -6,11 +6,11 @@ import pytest from areal.api.cli_args import SchedulingSpec, TrainEngineConfig -from areal.experimental.training_service.controller.controller import ( +from areal.v2.training_service.controller.controller import ( GatewayTrainController, ) -MODULE = "areal.experimental.training_service.controller.controller" +MODULE = "areal.v2.training_service.controller.controller" def _make_response(method: str, url: str, *, json=None) -> httpx.Response: diff --git a/tests/experimental/training_service/test_data_proxy_unit.py b/tests/v2/training_service/test_data_proxy_unit.py similarity index 98% rename from tests/experimental/training_service/test_data_proxy_unit.py rename to tests/v2/training_service/test_data_proxy_unit.py index 5a209dbdd9..2be179225d 100644 --- a/tests/experimental/training_service/test_data_proxy_unit.py +++ b/tests/v2/training_service/test_data_proxy_unit.py @@ -10,15 +10,15 @@ import pytest_asyncio import torch -from areal.experimental.training_service.data_proxy.app import create_app -from areal.experimental.training_service.data_proxy.config import TrainDataProxyConfig -from areal.experimental.training_service.data_proxy.dispatcher import Dispatcher -from areal.experimental.training_service.data_proxy.topology import ( +from areal.infra.controller.train_controller import _dispatch_tensors +from areal.infra.rpc.serialization import deserialize_value, serialize_value +from areal.v2.training_service.data_proxy.app import create_app +from areal.v2.training_service.data_proxy.config import TrainDataProxyConfig +from areal.v2.training_service.data_proxy.dispatcher import Dispatcher +from areal.v2.training_service.data_proxy.topology import ( WorkerInfo, WorkerTopology, ) -from areal.infra.controller.train_controller import _dispatch_tensors -from areal.infra.rpc.serialization import deserialize_value, serialize_value ADMIN_KEY = "dp-admin-key" diff --git a/tests/experimental/training_service/test_gateway_unit.py b/tests/v2/training_service/test_gateway_unit.py similarity index 97% rename from tests/experimental/training_service/test_gateway_unit.py rename to tests/v2/training_service/test_gateway_unit.py index 84c55dbcc7..0e314e0156 100644 --- a/tests/experimental/training_service/test_gateway_unit.py +++ b/tests/v2/training_service/test_gateway_unit.py @@ -8,16 +8,16 @@ import pytest import pytest_asyncio -from areal.experimental.training_service.gateway.app import create_app -from areal.experimental.training_service.gateway.config import GatewayConfig -from areal.experimental.training_service.gateway.streaming import ( +from areal.v2.training_service.gateway.app import create_app +from areal.v2.training_service.gateway.config import GatewayConfig +from areal.v2.training_service.gateway.streaming import ( RouterKeyRejectedError, RouterUnreachableError, forward_request, query_router, ) -MODULE = "areal.experimental.training_service.gateway.app" +MODULE = "areal.v2.training_service.gateway.app" ADMIN_KEY = "test-admin-key" SESSION_KEY = "session-key" WORKER_ADDR = "http://mock-worker:18082" diff --git a/tests/experimental/training_service/test_router_unit.py b/tests/v2/training_service/test_router_unit.py similarity index 97% rename from tests/experimental/training_service/test_router_unit.py rename to tests/v2/training_service/test_router_unit.py index 2f9bae4295..ad172d465d 100644 --- a/tests/experimental/training_service/test_router_unit.py +++ b/tests/v2/training_service/test_router_unit.py @@ -8,11 +8,11 @@ import pytest import pytest_asyncio -from areal.experimental.training_service.router.app import ( +from areal.v2.training_service.router.app import ( _probe_model_health, create_app, ) -from areal.experimental.training_service.router.config import RouterConfig +from areal.v2.training_service.router.config import RouterConfig ADMIN_KEY = "router-admin-key" diff --git a/tests/experimental/training_service/test_worker_unit.py b/tests/v2/training_service/test_worker_unit.py similarity index 87% rename from tests/experimental/training_service/test_worker_unit.py rename to tests/v2/training_service/test_worker_unit.py index 8d3c71e7aa..69b080625c 100644 --- a/tests/experimental/training_service/test_worker_unit.py +++ b/tests/v2/training_service/test_worker_unit.py @@ -6,15 +6,15 @@ import pytest -from areal.experimental.training_service.worker.config import TrainWorkerConfig from areal.infra.rpc.serialization import deserialize_value, serialize_value +from areal.v2.training_service.worker.config import TrainWorkerConfig -MODULE = "areal.experimental.training_service.worker.app" +MODULE = "areal.v2.training_service.worker.app" @pytest.fixture(autouse=True) def reset_worker_state(): - import areal.experimental.training_service.worker.app as worker_app + import areal.v2.training_service.worker.app as worker_app if worker_app._engine_work_queue is not None: worker_app._engine_work_queue.put(None) @@ -41,7 +41,7 @@ def reset_worker_state(): @pytest.fixture def client(): - from areal.experimental.training_service.worker.app import create_app + from areal.v2.training_service.worker.app import create_app app = create_app( TrainWorkerConfig( @@ -66,7 +66,7 @@ def test_create_engine_success(self, client): resp = client.post( "/create_engine", json={ - "engine_class": "tests.experimental.training_service.fake_train_engine.FakeTrainEngine", + "engine_class": "tests.v2.training_service.fake_train_engine.FakeTrainEngine", "init_args": serialize_value([]), "init_kwargs": serialize_value({"world_size": 1}), }, @@ -86,7 +86,7 @@ def test_train_batch_after_create_engine(self, client): create_resp = client.post( "/create_engine", json={ - "engine_class": "tests.experimental.training_service.fake_train_engine.FakeTrainEngine", + "engine_class": "tests.v2.training_service.fake_train_engine.FakeTrainEngine", "init_args": serialize_value([]), "init_kwargs": serialize_value({"world_size": 1}), }, @@ -111,7 +111,7 @@ def test_topology_after_create_engine(self, client): create_resp = client.post( "/create_engine", json={ - "engine_class": "tests.experimental.training_service.fake_train_engine.FakeTrainEngine", + "engine_class": "tests.v2.training_service.fake_train_engine.FakeTrainEngine", "init_args": serialize_value([]), "init_kwargs": serialize_value({"world_size": 1}), }, @@ -134,7 +134,7 @@ def test_ppo_endpoints_return_400_when_engine_method_missing(self, client): create_resp = client.post( "/create_engine", json={ - "engine_class": "tests.experimental.training_service.fake_train_engine.FakeTrainEngine", + "engine_class": "tests.v2.training_service.fake_train_engine.FakeTrainEngine", "init_args": serialize_value([]), "init_kwargs": serialize_value({"world_size": 1}), }, @@ -162,7 +162,7 @@ def test_forward_batch_after_initialize_succeeds_without_distributed_group( create_resp = client.post( "/create_engine", json={ - "engine_class": "tests.experimental.training_service.fake_train_engine.FakeTrainEngine", + "engine_class": "tests.v2.training_service.fake_train_engine.FakeTrainEngine", "init_args": serialize_value([]), "init_kwargs": serialize_value({"world_size": 1}), }, @@ -198,7 +198,7 @@ def test_sft_route_succeeds_without_distributed_group_for_single_worker( create_resp = client.post( "/create_engine", json={ - "engine_class": "tests.experimental.training_service.fake_train_engine.FakeTrainEngine", + "engine_class": "tests.v2.training_service.fake_train_engine.FakeTrainEngine", "init_args": serialize_value([]), "init_kwargs": serialize_value({"world_size": 1}), }, @@ -221,7 +221,7 @@ def test_sft_route_ignores_rpc_meta_override_for_single_worker(self, client): create_resp = client.post( "/create_engine", json={ - "engine_class": "tests.experimental.training_service.fake_train_engine.FakeTrainEngine", + "engine_class": "tests.v2.training_service.fake_train_engine.FakeTrainEngine", "init_args": serialize_value([]), "init_kwargs": serialize_value({"world_size": 1}), }, diff --git a/tests/v2/weight_update/__init__.py b/tests/v2/weight_update/__init__.py new file mode 100644 index 0000000000..9881313609 --- /dev/null +++ b/tests/v2/weight_update/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/experimental/weight_update/conftest.py b/tests/v2/weight_update/conftest.py similarity index 86% rename from tests/experimental/weight_update/conftest.py rename to tests/v2/weight_update/conftest.py index 56ed08be50..0223d78345 100644 --- a/tests/experimental/weight_update/conftest.py +++ b/tests/v2/weight_update/conftest.py @@ -3,7 +3,7 @@ The top-level ``areal/__init__.py`` imports from modules that use PEP 695 syntax (``def func[T](...)``) which is only valid on Python 3.12+. When running tests on Python 3.10/3.11 we register lightweight stubs for the ``areal`` namespace packages -so that importing ``areal.experimental.weight_update.*`` never triggers the +so that importing ``areal.v2.weight_update.*`` never triggers the problematic top-level init. """ @@ -18,9 +18,10 @@ _STUB_PACKAGES = [ ("areal", os.path.join(_REPO_ROOT, "areal")), ("areal.experimental", os.path.join(_REPO_ROOT, "areal", "experimental")), + ("areal.v2", os.path.join(_REPO_ROOT, "areal", "v2")), ( - "areal.experimental.weight_update", - os.path.join(_REPO_ROOT, "areal", "experimental", "weight_update"), + "areal.v2.weight_update", + os.path.join(_REPO_ROOT, "areal", "v2", "weight_update"), ), ] diff --git a/tests/experimental/weight_update/test_disk_integration.py b/tests/v2/weight_update/test_disk_integration.py similarity index 96% rename from tests/experimental/weight_update/test_disk_integration.py rename to tests/v2/weight_update/test_disk_integration.py index 7f400526f6..b5b7f2d225 100644 --- a/tests/experimental/weight_update/test_disk_integration.py +++ b/tests/v2/weight_update/test_disk_integration.py @@ -9,18 +9,18 @@ import pytest import torch -from tests.experimental.weight_update.test_nccl_integration import ( +from tests.v2.weight_update.test_nccl_integration import ( _get_test_model_path, _make_local_scheduler, _validate_weight_update_correctness, ) -from areal.experimental.weight_update.gateway.app import create_app -from areal.experimental.weight_update.gateway.config import ( +from areal.infra.platforms import current_platform +from areal.v2.weight_update.gateway.app import create_app +from areal.v2.weight_update.gateway.config import ( PairInfo, WeightUpdateConfig, ) -from areal.infra.platforms import current_platform @pytest.fixture() @@ -308,13 +308,13 @@ def test_disk_e2e_weight_update(n_gpus, tmp_path_factory): SchedulingSpec, TrainEngineConfig, ) - from areal.experimental.inference_service.controller.controller import ( + from areal.v2.inference_service.controller.controller import ( RolloutControllerV2, ) - from areal.experimental.training_service.controller.controller import ( + from areal.v2.training_service.controller.controller import ( GatewayTrainController, ) - from areal.experimental.weight_update.controller import ( + from areal.v2.weight_update.controller import ( WeightUpdateController, WeightUpdateControllerConfig, ) @@ -332,7 +332,7 @@ def test_disk_e2e_weight_update(n_gpus, tmp_path_factory): scheduling_spec=( SchedulingSpec( gpu=1, - cmd="python -m areal.experimental.inference_service.guard", + cmd="python -m areal.v2.inference_service.guard", ), ), consumer_batch_size=8, @@ -353,7 +353,7 @@ def test_disk_e2e_weight_update(n_gpus, tmp_path_factory): scheduling_spec=( SchedulingSpec( gpu=1, - cmd="python -m areal.experimental.training_service.guard", + cmd="python -m areal.v2.training_service.guard", env_vars=dict(NCCL_CUMEM_ENABLE="0", NCCL_NVLS_ENABLE="0"), ), ), diff --git a/tests/experimental/weight_update/test_kv_store.py b/tests/v2/weight_update/test_kv_store.py similarity index 97% rename from tests/experimental/weight_update/test_kv_store.py rename to tests/v2/weight_update/test_kv_store.py index 1e1f1576f8..f3f61b4ac4 100644 --- a/tests/experimental/weight_update/test_kv_store.py +++ b/tests/v2/weight_update/test_kv_store.py @@ -5,7 +5,7 @@ import pytest -from areal.experimental.weight_update.gateway.kv_store import WeightMetaStore +from areal.v2.weight_update.gateway.kv_store import WeightMetaStore @pytest.fixture() diff --git a/tests/experimental/weight_update/test_nccl_integration.py b/tests/v2/weight_update/test_nccl_integration.py similarity index 96% rename from tests/experimental/weight_update/test_nccl_integration.py rename to tests/v2/weight_update/test_nccl_integration.py index 8b9751b9fb..481df2c2bb 100644 --- a/tests/experimental/weight_update/test_nccl_integration.py +++ b/tests/v2/weight_update/test_nccl_integration.py @@ -34,7 +34,7 @@ def _run_weight_update_test(n_gpus: int, test_type: str, output: str): "--nnodes=1", "--master-addr=localhost", f"--master_port={port}", - "tests/experimental/weight_update/torchrun/run_nccl_weight_transfer.py", + "tests/v2/weight_update/torchrun/run_nccl_weight_transfer.py", f"--test_type={test_type}", f"--output={output}", ], @@ -374,13 +374,13 @@ def test_awex_fsdp_e2e_weight_update(n_gpus, tmp_path_factory): SchedulingSpec, TrainEngineConfig, ) - from areal.experimental.inference_service.controller.controller import ( + from areal.v2.inference_service.controller.controller import ( RolloutControllerV2, ) - from areal.experimental.training_service.controller.controller import ( + from areal.v2.training_service.controller.controller import ( GatewayTrainController, ) - from areal.experimental.weight_update.controller import ( + from areal.v2.weight_update.controller import ( WeightUpdateController, WeightUpdateControllerConfig, ) @@ -397,7 +397,7 @@ def test_awex_fsdp_e2e_weight_update(n_gpus, tmp_path_factory): scheduling_spec=( SchedulingSpec( gpu=1, - cmd="python -m areal.experimental.inference_service.guard", + cmd="python -m areal.v2.inference_service.guard", ), ), consumer_batch_size=8, @@ -418,7 +418,7 @@ def test_awex_fsdp_e2e_weight_update(n_gpus, tmp_path_factory): scheduling_spec=( SchedulingSpec( gpu=1, - cmd="python -m areal.experimental.training_service.guard", + cmd="python -m areal.v2.training_service.guard", env_vars=dict(NCCL_CUMEM_ENABLE="0", NCCL_NVLS_ENABLE="0"), ), ), @@ -577,13 +577,13 @@ def _run_megatron_awex_e2e( SchedulingSpec, TrainEngineConfig, ) - from areal.experimental.inference_service.controller.controller import ( + from areal.v2.inference_service.controller.controller import ( RolloutControllerV2, ) - from areal.experimental.training_service.controller.controller import ( + from areal.v2.training_service.controller.controller import ( GatewayTrainController, ) - from areal.experimental.weight_update.controller import ( + from areal.v2.weight_update.controller import ( WeightUpdateController, WeightUpdateControllerConfig, ) @@ -599,7 +599,7 @@ def _run_megatron_awex_e2e( scheduling_spec=( SchedulingSpec( gpu=1, - cmd="python -m areal.experimental.inference_service.guard", + cmd="python -m areal.v2.inference_service.guard", ), ), consumer_batch_size=8, @@ -621,7 +621,7 @@ def _run_megatron_awex_e2e( scheduling_spec=( SchedulingSpec( gpu=1, - cmd="python -m areal.experimental.training_service.guard", + cmd="python -m areal.v2.training_service.guard", env_vars=dict(NCCL_CUMEM_ENABLE="0", NCCL_NVLS_ENABLE="0"), ), ), @@ -923,13 +923,13 @@ def _run_megatron_colocate_e2e( SchedulingSpec, TrainEngineConfig, ) - from areal.experimental.inference_service.controller.controller import ( + from areal.v2.inference_service.controller.controller import ( RolloutControllerV2, ) - from areal.experimental.training_service.controller.controller import ( + from areal.v2.training_service.controller.controller import ( GatewayTrainController, ) - from areal.experimental.weight_update.controller import ( + from areal.v2.weight_update.controller import ( WeightUpdateController, WeightUpdateControllerConfig, ) @@ -945,7 +945,7 @@ def _run_megatron_colocate_e2e( scheduling_spec=( SchedulingSpec( gpu=1, - cmd="python -m areal.experimental.inference_service.guard", + cmd="python -m areal.v2.inference_service.guard", ), ), consumer_batch_size=8, @@ -966,7 +966,7 @@ def _run_megatron_colocate_e2e( scheduling_spec=( SchedulingSpec( gpu=1, - cmd="python -m areal.experimental.training_service.guard", + cmd="python -m areal.v2.training_service.guard", env_vars=dict(NCCL_CUMEM_ENABLE="0", NCCL_NVLS_ENABLE="0"), ), ), @@ -1093,13 +1093,13 @@ def test_awex_megatron_colocate_dp_multi_version_e2e(tmp_path_factory): SchedulingSpec, TrainEngineConfig, ) - from areal.experimental.inference_service.controller.controller import ( + from areal.v2.inference_service.controller.controller import ( RolloutControllerV2, ) - from areal.experimental.training_service.controller.controller import ( + from areal.v2.training_service.controller.controller import ( GatewayTrainController, ) - from areal.experimental.weight_update.controller import ( + from areal.v2.weight_update.controller import ( WeightUpdateController, WeightUpdateControllerConfig, ) @@ -1115,7 +1115,7 @@ def test_awex_megatron_colocate_dp_multi_version_e2e(tmp_path_factory): scheduling_spec=( SchedulingSpec( gpu=1, - cmd="python -m areal.experimental.inference_service.guard", + cmd="python -m areal.v2.inference_service.guard", ), ), consumer_batch_size=8, @@ -1136,7 +1136,7 @@ def test_awex_megatron_colocate_dp_multi_version_e2e(tmp_path_factory): scheduling_spec=( SchedulingSpec( gpu=1, - cmd="python -m areal.experimental.training_service.guard", + cmd="python -m areal.v2.training_service.guard", env_vars=dict(NCCL_CUMEM_ENABLE="0", NCCL_NVLS_ENABLE="0"), ), ), diff --git a/tests/experimental/weight_update/test_pair_registry.py b/tests/v2/weight_update/test_pair_registry.py similarity index 94% rename from tests/experimental/weight_update/test_pair_registry.py rename to tests/v2/weight_update/test_pair_registry.py index 89b9009250..dc4fd56ff5 100644 --- a/tests/experimental/weight_update/test_pair_registry.py +++ b/tests/v2/weight_update/test_pair_registry.py @@ -3,8 +3,8 @@ import pytest -from areal.experimental.weight_update.gateway.config import PairInfo -from areal.experimental.weight_update.gateway.pair_registry import PairRegistry +from areal.v2.weight_update.gateway.config import PairInfo +from areal.v2.weight_update.gateway.pair_registry import PairRegistry @pytest.fixture diff --git a/tests/experimental/weight_update/test_sglang_integration.py b/tests/v2/weight_update/test_sglang_integration.py similarity index 96% rename from tests/experimental/weight_update/test_sglang_integration.py rename to tests/v2/weight_update/test_sglang_integration.py index 0cf1b19e92..e5851950a3 100644 --- a/tests/experimental/weight_update/test_sglang_integration.py +++ b/tests/v2/weight_update/test_sglang_integration.py @@ -4,7 +4,7 @@ Requires GPU. Marked @pytest.mark.slow and @pytest.mark.sglang to exclude from default CI. Run manually: - uv run pytest tests/experimental/weight_update/test_sglang_server_integration.py -v -s + uv run pytest tests/v2/weight_update/test_sglang_server_integration.py -v -s """ from __future__ import annotations @@ -52,7 +52,7 @@ def sglang_server(): [ "python", "-m", - "areal.experimental.inference_service.sglang.launch_server", + "areal.v2.inference_service.sglang.launch_server", "--model-path", model_path, "--port", diff --git a/tests/experimental/weight_update/test_wu_controller.py b/tests/v2/weight_update/test_wu_controller.py similarity index 95% rename from tests/experimental/weight_update/test_wu_controller.py rename to tests/v2/weight_update/test_wu_controller.py index f1aca5531e..68e3b93432 100644 --- a/tests/experimental/weight_update/test_wu_controller.py +++ b/tests/v2/weight_update/test_wu_controller.py @@ -7,13 +7,13 @@ import pytest import requests -from areal.experimental.weight_update.controller.config import ( +from areal.v2.weight_update.controller.config import ( WeightUpdateControllerConfig, ) -from areal.experimental.weight_update.controller.controller import ( +from areal.v2.weight_update.controller.controller import ( WeightUpdateController, ) -from areal.experimental.weight_update.gateway.config import WeightUpdateResult +from areal.v2.weight_update.gateway.config import WeightUpdateResult GATEWAY_URL = "http://localhost:7080" @@ -81,6 +81,7 @@ def test_connect_sends_correct_request(self, ctrl): "save_path": "", "use_lora": False, "lora_name": "", + "lora_keep_versions": 0, "colocate": False, "nccl_master_addr": "", "nccl_master_port": 0, @@ -113,6 +114,7 @@ def test_connect_disk_mode_sends_disk_fields(self, ctrl): "save_path": "/shared/weights", "use_lora": True, "lora_name": "my-lora", + "lora_keep_versions": 0, "colocate": False, "nccl_master_addr": "", "nccl_master_port": 0, diff --git a/tests/v2/weight_update/torchrun/__init__.py b/tests/v2/weight_update/torchrun/__init__.py new file mode 100644 index 0000000000..9881313609 --- /dev/null +++ b/tests/v2/weight_update/torchrun/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/experimental/weight_update/torchrun/dist_utils.py b/tests/v2/weight_update/torchrun/dist_utils.py similarity index 100% rename from tests/experimental/weight_update/torchrun/dist_utils.py rename to tests/v2/weight_update/torchrun/dist_utils.py diff --git a/tests/experimental/weight_update/torchrun/run_nccl_weight_transfer.py b/tests/v2/weight_update/torchrun/run_nccl_weight_transfer.py similarity index 98% rename from tests/experimental/weight_update/torchrun/run_nccl_weight_transfer.py rename to tests/v2/weight_update/torchrun/run_nccl_weight_transfer.py index 7931988f27..abfcab0336 100644 --- a/tests/experimental/weight_update/torchrun/run_nccl_weight_transfer.py +++ b/tests/v2/weight_update/torchrun/run_nccl_weight_transfer.py @@ -7,16 +7,16 @@ import torch import torch.distributed as dist -from tests.experimental.weight_update.torchrun.dist_utils import ( +from tests.v2.weight_update.torchrun.dist_utils import ( print_rank0, write_result, ) -from areal.experimental.weight_update.nccl_group import ( +from areal.infra.platforms import current_platform +from areal.v2.weight_update.nccl_group import ( init_weights_update_group, setup_batch_isend_irecv, ) -from areal.infra.platforms import current_platform def run_nccl_group_init(output=None): diff --git a/uv.lock b/uv.lock index e468608cc3..3ee372d5e1 100644 --- a/uv.lock +++ b/uv.lock @@ -308,7 +308,7 @@ wheels = [ [[package]] name = "areal" -version = "1.0.4" +version = "2.0.0" source = { editable = "." } dependencies = [ { name = "aiofiles", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -400,6 +400,10 @@ dependencies = [ ] [package.optional-dependencies] +cli = [ + { name = "click", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "colorlog", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] cuda = [ { name = "flash-linear-attention", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "kernels", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -480,8 +484,10 @@ requires-dist = [ { name = "blosc" }, { name = "build", specifier = ">=1.2.1" }, { name = "claude-agent-sdk" }, + { name = "click", marker = "extra == 'cli'", specifier = ">=8.1" }, { name = "colorama" }, { name = "colorlog" }, + { name = "colorlog", marker = "extra == 'cli'" }, { name = "cookiecutter", specifier = ">2.1.1" }, { name = "datasets", specifier = ">=3.0.0" }, { name = "daytona", marker = "extra == 'sandbox'", specifier = ">=0.167.0" }, @@ -573,7 +579,7 @@ requires-dist = [ { name = "word2number" }, { name = "zstandard" }, ] -provides-extras = ["sglang", "tms", "kernels", "megatron", "cuda-train", "cuda", "sandbox"] +provides-extras = ["sglang", "tms", "kernels", "megatron", "cuda-train", "cuda", "sandbox", "cli"] [package.metadata.requires-dev] dev = [ diff --git a/uv.vllm.lock b/uv.vllm.lock index 9e9b587e61..5cbac42895 100644 --- a/uv.vllm.lock +++ b/uv.vllm.lock @@ -347,7 +347,7 @@ wheels = [ [[package]] name = "areal" -version = "1.0.4" +version = "2.0.0" source = { editable = "." } dependencies = [ { name = "aiofiles", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -448,6 +448,10 @@ dependencies = [ ] [package.optional-dependencies] +cli = [ + { name = "click", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "colorlog", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, +] cuda = [ { name = "flash-linear-attention", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "kernels", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'x86_64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -526,8 +530,10 @@ requires-dist = [ { name = "blosc" }, { name = "build", specifier = ">=1.2.1" }, { name = "claude-agent-sdk" }, + { name = "click", marker = "extra == 'cli'", specifier = ">=8.1" }, { name = "colorama" }, { name = "colorlog" }, + { name = "colorlog", marker = "extra == 'cli'" }, { name = "cookiecutter", specifier = ">2.1.1" }, { name = "datasets", specifier = ">=3.0.0" }, { name = "daytona", marker = "extra == 'sandbox'", specifier = ">=0.167.0" }, @@ -619,7 +625,7 @@ requires-dist = [ { name = "word2number" }, { name = "zstandard" }, ] -provides-extras = ["vllm", "tms", "kernels", "megatron", "cuda-train", "cuda", "sandbox"] +provides-extras = ["vllm", "tms", "kernels", "megatron", "cuda-train", "cuda", "sandbox", "cli"] [package.metadata.requires-dev] dev = [