Setting up Multi-Node Ray or Slurm Clusters for Large-Scale Training with AReaL

AReaL provides production-ready RayScheduler and SlurmScheduler classes that abstract resource allocation, placement-group management, and worker lifecycle behind a unified Python API for distributed training.

Setting up multi-node Ray or Slurm clusters for large-scale training requires handling complex orchestration logic, from GPU placement to worker discovery. The AReaL repository (inclusionai/areal) ships with two cluster-launchers—RayScheduler and SlurmScheduler—that implement the common abstract base areal.infra.scheduler.Scheduler. Both schedulers share a worker abstraction (areal.infra.scheduler.Worker) for network endpoints and a scheduling spec (areal.api.cli_args.SchedulingSpec) for resource requirements, enabling seamless switching between elastic cloud clusters and traditional HPC environments.

RayScheduler: Elastic GPU Sharing with Placement Groups

The RayScheduler (implemented in areal/infra/scheduler/ray.py) targets dynamic environments such as Kubernetes or cloud VMs running Ray 2.x. It uses placement groups to enforce resource isolation while supporting elastic scaling.

Resource Detection and Bundle Creation

Before launching workers, ray_resource_type() (lines 40-49 of ray.py) inspects the local environment using torch.cuda.is_available() and NPU detection to report "GPU", "NPU", or "CPU". For each worker, the scheduler creates a dedicated placement group with a single bundle matching the exact resource requirements (_bundle_spec, _create_placement_group). This guarantees exclusive GPU access for the actor.

Actor Launch and RPC Server

Workers are instantiated as Ray actors of RayRPCServer (lines 85-90). The runtime_env is populated with environment variables from get_env_vars and get_thread_env_vars. A PlacementGroupSchedulingStrategy forces the actor onto the pre-created group (see strategy_kwargs lines 79-83).

Master IP and Port Discovery

The first placement group’s master address is obtained via get_placement_group_master_ip_and_port (areal/infra/utils/ray.py lines 8-15). This spawns a tiny Ray actor that binds to a free port on the node’s host IP, ensuring workers can locate the master without manual configuration.

Health Checks and Colocation

The _ping_workers method (lines 92-105) repeatedly calls a ping method on each actor and raises WorkerTimeoutError if any actor fails to respond. For colocation without forking, the target role’s existing workers are reused. For forked colocation, _create_forked_workers_internal creates new actors sharing the same placement group and bundle index, using a reduced GPU allocation (0.01 GPU) so multiple forked actors coexist on a single GPU (lines 60-68).

SlurmScheduler: Traditional HPC with sbatch Integration

The SlurmScheduler (implemented in areal/infra/scheduler/slurm.py) targets conventional HPC clusters where jobs are submitted via sbatch and run under a shared filesystem.

Job Submission and Script Generation

The _generate_sbatch_script method (lines 322-384) builds a single-command sbatch script that runs one srun per node. The script injects the RPC server command, sets CUDA_VISIBLE_DEVICES based on SLURM_LOCALID, and streams logs to per-role files (_log_path_of, _merged_log_path). The scheduler computes the number of nodes required from the total GPU count (nodes = max(1, (total_gpus + n_gpus_per_node - 1) // n_gpus_per_node)) and passes that to sbatch.

Worker Discovery via Name Resolution

After submission, each worker calls the central name-resolve service (name_resolve) to publish its <ip>:<port> address. The controller repeatedly polls _discover_worker_network (lines 140-166) until all workers are marked discovered. This decouples discovery from Slurm’s job ID, allowing flexible network topologies.

Health Checks and Process Forking

The _is_worker_ready method performs a simple HTTP GET /health request on the worker’s RPC server (lines 68-80). For forked colocation, _fork_single_worker sends an HTTP POST /fork to the parent worker; the parent spawns a new process on the same node, and the forked worker’s IP/port are returned. Cleanup is performed via /kill_forked_worker (lines 530-562). Forked workers have a sentinel slurm_job_id = -1 because they are not separate Slurm jobs.

Choosing Between Ray and Slurm for Large-Scale Training

Factor	RayScheduler	SlurmScheduler
Cluster type	Dynamic cloud VMs, Kubernetes, or on-prem clusters with Ray installed.	Traditional HPC clusters with Slurm installed and a shared filesystem.
Elasticity	Workers can be added/removed at runtime; placement groups are created per actor.	Jobs are static after sbatch – adding workers requires a new Slurm job.
Container support	Implicit via Ray-Docker or native; can also launch containerized RPC server manually.	Built-in container support (apptainer) via the container_type field in SchedulingSpec.
Port management	Ray finds a free port automatically (find_free_ports).	Slurm scripts compute a free port per node via a Bash seq / ss trick.
Debugging	Use Ray dashboard (ray status, ray logs).	Inspect sbatch output and Slurm logs (squeue, scontrol).
Typical scale	Up to a few hundred GPUs across elastic nodes.	Hundreds to thousands of GPUs on a fixed allocation.

Both schedulers expose the same high-level API, allowing you to switch backends by changing a single configuration field.

Code Examples

Launching a Minimal Ray Cluster

import ray
from areal.infra.scheduler.ray import RayScheduler
from areal.api.cli_args import BaseExperimentConfig, SchedulingSpec, SchedulingStrategy

# Initialise Ray (local mode for testing)

ray.init(ignore_reinit_error=True)

# Minimal experiment config

exp_cfg = BaseExperimentConfig(
    experiment_name="demo",
    trial_name="run1",
    cluster=type('C', (), {"n_gpus_per_node": 4})  # dummy cluster object

)

scheduler = RayScheduler(exp_config=exp_cfg)

# Create 2 workers, each with 1 GPU

job = type(
    "Job",
    (),
    {
        "role": "rollout",
        "replicas": 2,
        "tasks": [SchedulingSpec(cpu=4, gpu=1, mem=16, port_count=2)],
        "scheduling_strategy": SchedulingStrategy(type="separation")
    },
)()
worker_ids = scheduler.create_workers(job)
print("Ray workers:", worker_ids)

# Clean up

scheduler.delete_workers()
ray.shutdown()

Submitting a Slurm Job with Colocation

# On the login node, create submit.py

cat > submit.py <<'PY'
from areal.infra.scheduler.slurm import SlurmScheduler
from areal.api.cli_args import BaseExperimentConfig, SchedulingSpec, SchedulingStrategy

exp_cfg = BaseExperimentConfig(
    experiment_name="demo",
    trial_name="run1",
    cluster=type('C', (), {
        "n_gpus_per_node": 8,
        "fileroot": "/shared/areal_logs",
        "cluster_name": "mycluster"
    })
)

scheduler = SlurmScheduler(exp_config=exp_cfg)

# Primary rollout workers

rollout_job = type(
    "Job",
    (),
    {
        "role": "rollout",
        "replicas": 2,
        "tasks": [SchedulingSpec(cpu=8, gpu=1, mem=32)],
        "scheduling_strategy": SchedulingStrategy(type="separation")
    },
)()
rollout_ids = scheduler.create_workers(rollout_job)

# Fork proxy workers on the same nodes

proxy_ids = scheduler.fork_workers(
    role="proxy",
    target_role="rollout",
    command=None
)

print("Rollout IDs:", rollout_ids)
print("Proxy IDs:", proxy_ids)
PY

python submit.py

Integrating Schedulers into Training Scripts

from areal.infra import LocalScheduler, RayScheduler, SlurmScheduler
from areal.api.cli_args import load_expr_config, GRPOConfig

def main():
    # Load experiment configuration

    cfg, _ = load_expr_config(["examples/configs/gsm8k.yaml"], GRPOConfig)

    # Instantiate appropriate scheduler

    if cfg.scheduler.type == "ray":
        scheduler = RayScheduler(exp_config=cfg)
    elif cfg.scheduler.type == "slurm":
        scheduler = SlurmScheduler(exp_config=cfg)
    else:
        scheduler = LocalScheduler(exp_config=cfg)

    # Create rollout workers

    rollout_ids = scheduler.create_workers(
        job=type(
            "Job",
            (),
            {
                "role": "rollout",
                "replicas": cfg.rollout.n_replicas,
                "tasks": cfg.scheduling_spec,
                "scheduling_strategy": cfg.scheduling_strategy,
            },
        )()
    )

    # Invoke engine method remotely

    result = scheduler.call_engine(
        worker_id=rollout_ids[0],
        method="generate",
        engine_name="my_engine",
        inputs={"prompt": "Explain quantum computing"}
    )
    print(result)

    # Cleanup

    scheduler.delete_workers()

Summary

• AReaL provides unified abstractions for setting up multi-node Ray or Slurm clusters through RayScheduler and SlurmScheduler, both implementing the areal.infra.scheduler.Scheduler interface.
• RayScheduler leverages placement groups for elastic GPU allocation and supports dynamic worker forking with fractional GPU sharing (0.01 GPU per forked worker).
• SlurmScheduler generates sbatch scripts for traditional HPC environments, uses a name-resolve service for worker discovery, and supports process forking via HTTP /fork endpoints.
• Both schedulers share identical APIs for create_workers, call_engine, and delete_workers, allowing seamless switching between cloud elasticity and fixed HPC allocations by changing a configuration field.

Frequently Asked Questions

What is the difference between RayScheduler and SlurmScheduler in AReaL?

RayScheduler (in areal/infra/scheduler/ray.py) targets dynamic environments like Kubernetes or cloud VMs, using Ray placement groups to allocate GPUs elastically at runtime. SlurmScheduler (in areal/infra/scheduler/slurm.py) targets traditional HPC clusters, submitting static sbatch jobs that reserve nodes for the entire job duration. While Ray supports adding workers mid-job, Slurm requires a new job submission to change the worker count.

How does AReaL handle worker discovery in Slurm clusters?

Instead of relying on Slurm environment variables alone, AReaL uses a name-resolve service where each worker publishes its <ip>:<port> address via name_resolve.put. The controller polls _discover_worker_network (lines 140-166 of slurm.py) until all workers are marked discovered, decoupling network topology from Slurm's job ID system and enabling flexible multi-node communication.

Can I run forked workers on the same GPU using AReaL?

Yes. Both schedulers support forked colocation, but the implementation differs. In RayScheduler, _create_forked_workers_internal creates new actors sharing the same placement group with a fractional GPU allocation of 0.01 GPU (lines 60-68 of ray.py), allowing multiple forked actors to coexist on a single GPU. In SlurmScheduler, _fork_single_worker sends an HTTP POST /fork to the parent process, which spawns a new process on the same node without requesting additional Slurm resources.

Which scheduler should I choose for elastic scaling in large-scale training?

Choose RayScheduler if you need elasticity—adding or removing workers at runtime across dynamic cloud VMs or Kubernetes. It supports up to a few hundred GPUs with automatic port management and placement groups. Choose SlurmScheduler if you operate on a fixed-allocation HPC cluster with hundreds to thousands of GPUs, require strict queue management via sbatch, or need built-in container support through Apptainer. Both expose identical Python APIs, so you can switch by changing the scheduler.type configuration field without modifying your training logic.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}