refactor(infra): move scheduler and rpc modules under areal/infra

Consolidate infrastructure-related modules under areal/infra/ for better
code organization:

- Move areal/scheduler/ to areal/infra/scheduler/
- Move areal/scheduler/rpc/ to areal/infra/rpc/
- Update all imports across the codebase
- Update example configs to use new import paths
- Add new areal/infra/rpc/__init__.py for cleaner exports

Key changes:
- Scheduler classes (LocalScheduler, RayScheduler, SlurmScheduler) now at
  areal.infra.scheduler
- RPC utilities (RTensor, serialization) now at areal.infra.rpc
- All existing functionality preserved with updated import paths

Author: garrett4wade

.claude/agents/launcher-scheduler-expert.md

@@ -22,7 +22,8 @@ allocation correctness.
 
 Use this agent **when requested** when:
 
-- **Code modifications**: User edits files in `areal/launcher/` or `areal/scheduler/`
+- **Code modifications**: User edits files in `areal/launcher/`, `areal/infra/rpc`, or
+  `areal/infra/scheduler/`
 - **Configuration changes**: User modifies `ClusterSpecConfig`, `SchedulerConfig`, or
   related dataclasses
 - **Deployment issues**: User encounters job launch failures, port conflicts, GPU
@@ -99,8 +100,8 @@ Critical utilities in `areal/utils/launcher.py`:
   hard-coded GPU indices or direct `torch.cuda` calls
 - Use `areal.utils.name_resolve` for multi-node service discovery -> not direct
   IP/hostname assumptions
-- Raise specific exceptions from `areal.scheduler.exceptions` -> not generic exception
-  types
+- Raise specific exceptions from `areal.infra.scheduler.exceptions` -> not generic
+  exception types
 - Use `areal.utils.proc.kill_process_tree()` for process termination -> not leaving
   zombie processes
 - Propagate all `BASE_ENVIRONS` variables and thread control variables -> not missing
@@ -139,9 +140,9 @@ Critical utilities in `areal/utils/launcher.py`:
 | `areal/launcher/ray.py`           | Ray cluster deployment             | Ray actor management, placement group allocation          |
 | `areal/launcher/sglang_server.py` | SGLang inference server deployment | SGLang server process management, cache isolation         |
 | `areal/launcher/vllm_server.py`   | vLLM inference server deployment   | vLLM server process management, cache isolation           |
-| `areal/scheduler/local.py`        | Local worker scheduling            | GPU round-robin, port allocation, health monitoring       |
-| `areal/scheduler/slurm.py`        | Slurm-integrated scheduling        | Job array coordination, resource reservation              |
-| `areal/scheduler/ray.py`          | Ray cluster scheduling             | Ray placement groups, actor-based worker management       |
+| `areal/infra/scheduler/local.py`  | Local worker scheduling            | GPU round-robin, port allocation, health monitoring       |
+| `areal/infra/scheduler/slurm.py`  | Slurm-integrated scheduling        | Job array coordination, resource reservation              |
+| `areal/infra/scheduler/ray.py`    | Ray cluster scheduling             | Ray placement groups, actor-based worker management       |
 | `areal/utils/launcher.py`         | Shared utilities                   | Environment variable management, configuration validation |
 
 ______________________________________________________________________

.claude/data/pr-review-change-types.md

@@ -44,7 +44,7 @@ ______________________________________________________________________
 | **CHECKPOINT_RECOVERY** | `areal/utils/saver.py`, `areal/utils/recover.py`, `areal/utils/fsdp/checkpoint.py` | `state_dict`, `load_state_dict`, `checkpoint`                            |
 | **REWARD**              | `areal/reward/`                                                                    | `reward_fn`, `AsyncRewardWrapper`, `MathVerifyWorker`                    |
 | **DATASET**             | `areal/dataset/`                                                                   | `get_*_dataset`, `DataLoader`, `IterableDataset`                         |
-| **LAUNCHER_SCHEDULER**  | `areal/launcher/`, `areal/scheduler/`                                              | `LaunchConfig`, `Scheduler`, `RayLauncher`, `SlurmLauncher`              |
+| **LAUNCHER_SCHEDULER**  | `areal/launcher/`, `areal/infra/scheduler/`, `areal/infra/rpc`                     | `LaunchConfig`, `Scheduler`, `RayLauncher`, `SlurmLauncher`              |
 | **ATTENTION**           | `attention.py`, `varlen_attention.py`                                              | `flash_attn`, `sdpa`, `varlen`, `causal_mask`                            |
 
 ## LOW Level (Use Haiku)

.claude/data/pr-review-templates.md

@@ -349,7 +349,7 @@ Checklist:
 ### Launcher and Scheduler Configuration \[Sonnet\]
 
 ```
-Applicable: areal/launcher/, areal/scheduler/ directories
+Applicable: areal/launcher/, areal/infra/scheduler/, areal/infra/rpc directories
 Checklist:
 - Resource config reasonableness (GPU count, memory)
 - Process group config matches parallel strategy

AGENTS.md

@@ -37,8 +37,9 @@ When unsure, leave a `TODO(agent)` comment and note the constraint in your respo
   shared utilities:
   - `areal/api/` - Contracts for workflows, engines, schedulers, IO structs, and
     CLI/config dataclasses.
-  - `areal/infra/` - Single controller implementation, async orchestration primitives,
-    and hardware/platform abstractions for CPU/GPU/NPU runtimes.
+  - `areal/infra/` - Core infrastructure including single controller implementation,
+    placement and allocation policies, async orchestration primitives, and
+    hardware/platform abstractions for CPU/GPU/NPU runtimes.
   - `areal/dataset/` - Stateful dataset loaders (GSM8K, Geometry3K, CLEVR, HH-RLHF,
     TORL, etc.) and utilities that feed rollout jobs safely.
   - `areal/engine/` - Training backends (FSDP2, Megatron, PPO, SFT, reward modeling) and
@@ -51,8 +52,6 @@ When unsure, leave a `TODO(agent)` comment and note the constraint in your respo
     wrappers, custom heads).
   - `areal/reward/` - Built-in reward functions (GSM8K, Geometry3K, CLEVR, etc.), math
     parsers, and helpers; wrap slow logic with `AsyncRewardWrapper`.
-  - `areal/scheduler/` - Placement and allocation policies aligned with launcher
-    configs.
   - `areal/tests/` - Focused unit/integration suites (many require GPUs or mocked
     distributed backends).
   - `areal/tools/` - Developer utilities and maintenance scripts tied to the core

areal/__init__.py

@@ -3,22 +3,22 @@
 from .version import __version__  # noqa
 
 from .infra import (
-    TrainController,
     RolloutController,
-    WorkflowExecutor,
     StalenessManager,
-    workflow_context,
+    TrainController,
+    WorkflowExecutor,
     current_platform,
+    workflow_context,
 )
 from .trainer import PPOTrainer, SFTTrainer
 
 __all__ = [
-    "TrainController",
+    "PPOTrainer",
     "RolloutController",
-    "WorkflowExecutor",
+    "SFTTrainer",
     "StalenessManager",
-    "workflow_context",
+    "TrainController",
+    "WorkflowExecutor",
     "current_platform",
-    "PPOTrainer",
-    "SFTTrainer",
+    "workflow_context",
 ]

areal/api/cli_args.py

@@ -882,7 +882,7 @@ class TrainEngineConfig:
     # Scheduling
     scheduling_spec: tuple[SchedulingSpec, ...] = field(
         default_factory=lambda: (
-            SchedulingSpec(cmd="python -m areal.scheduler.rpc.rpc_server"),
+            SchedulingSpec(cmd="python -m areal.infra.rpc.rpc_server"),
         ),
         metadata={
             "help": "Train engine schedule specs. Can accept 1 or 2 SchedulingSpec: "
@@ -1512,7 +1512,7 @@ class InferenceEngineConfig:
     )
     scheduling_spec: tuple[SchedulingSpec, ...] = field(
         default_factory=lambda: (
-            SchedulingSpec(cmd="python -m areal.scheduler.rpc.rpc_server"),
+            SchedulingSpec(cmd="python -m areal.infra.rpc.rpc_server"),
         ),
         metadata={
             "help": "inference engine schedule specs. Can accept 1 or 2 SchedulingSpec: "

areal/experimental/openai/proxy/proxy_rollout_server.py

@@ -27,7 +27,7 @@
 
 from areal.api.cli_args import NameResolveConfig, OpenAIProxyConfig
 from areal.experimental.openai.client import ArealOpenAI
-from areal.scheduler.rpc.serialization import deserialize_value, serialize_value
+from areal.infra.rpc.serialization import deserialize_value, serialize_value
 from areal.utils import name_resolve, names, seeding
 from areal.utils.dynamic_import import import_from_string
 from areal.utils.hf_utils import load_hf_tokenizer

areal/infra/__init__.py

@@ -7,6 +7,7 @@
     RemoteInfBackendProtocol,
     RemoteInfEngine,
 )
+from .scheduler import LocalScheduler, RayScheduler, SlurmScheduler
 from .staleness_manager import StalenessManager
 from .workflow_executor import (
     WorkflowExecutor,
@@ -25,4 +26,7 @@
     "Platform",
     "current_platform",
     "is_npu_available",
+    "LocalScheduler",
+    "RayScheduler",
+    "SlurmScheduler",
 ]

areal/infra/controller/rollout_callback.py

@@ -5,7 +5,7 @@
 import requests
 
 from areal.api.io_struct import ParamSpec, WeightUpdateMeta
-from areal.scheduler.rpc.serialization import serialize_value
+from areal.infra.rpc.serialization import serialize_value
 from areal.utils import logging
 from areal.utils.concurrent import get_executor
 

areal/infra/controller/rollout_controller.py

@@ -29,7 +29,7 @@
 )
 from areal.api.scheduler_api import Job, Scheduler, Worker
 from areal.api.workflow_api import RolloutWorkflow, WorkflowLike
-from areal.scheduler.rpc.serialization import deserialize_value
+from areal.infra.rpc.serialization import deserialize_value
 from areal.utils import logging, perf_tracer
 from areal.utils.concurrent import run_async_task
 from areal.utils.data import concat_padded_tensors, cycle_dataloader