diff --git a/arealite/engine/sglang_remote.py b/arealite/engine/sglang_remote.py index c904d257e0..bb359b6dc2 100644 --- a/arealite/engine/sglang_remote.py +++ b/arealite/engine/sglang_remote.py @@ -44,6 +44,17 @@ class RemoteSGLangEngine(InferenceEngine): def __init__(self, config: InferenceEngineConfig): + """ + Initialize RemoteSGLangEngine that connects to pre-existing SGLang servers. + + IMPORTANT: This class expects SGLang servers to be ALREADY RUNNING. + In the AReaL training pipeline: + 1. Launcher starts SGLang servers first + 2. Launcher waits for servers to be ready + 3. Launcher starts training script with server addresses in AREAL_LLM_SERVER_ADDRS + 4. Training script creates RemoteSGLangEngine instances (this class) + 5. This class connects to the pre-existing servers + """ config.max_concurrent_rollouts = ( config.max_concurrent_rollouts or config.consumer_batch_size ) @@ -55,11 +66,16 @@ def __init__(self, config: InferenceEngineConfig): self.addresses = os.getenv("AREAL_LLM_SERVER_ADDRS").split(",") if not self.addresses: - raise RuntimeError("No configured SGLang servers.") + raise RuntimeError( + "No configured SGLang servers found in AREAL_LLM_SERVER_ADDRS. " + "SGLang servers must be started before creating RemoteSGLangEngine. " + "Use the AReaL launcher to start servers first." + ) + logger.info(f"RemoteSGLangEngine connecting to {len(self.addresses)} pre-existing SGLang servers...") logger.info("Waiting for server ready...") for addr in self.addresses: self._wait_for_server(addr) - logger.info("Servers are all ready!") + logger.info("All SGLang servers are ready and connected!") self.server_idx = random.randint(0, len(self.addresses) - 1) diff --git a/arealite/launcher/local.py b/arealite/launcher/local.py index bd73a216a3..fbc8aa2660 100644 --- a/arealite/launcher/local.py +++ b/arealite/launcher/local.py @@ -268,6 +268,11 @@ def main_local(): raise NotImplementedError() # Launch inference servers. + logger.info("=" * 60) + logger.info("STARTUP SEQUENCE: Starting SGLang inference servers FIRST") + logger.info("=" * 60) + logger.info(f"Launching {alloc_mode.gen_dp_size} SGLang server(s) with {alloc_mode.gen_tp_size} tensor parallel size each") + launcher.submit_array( job_name="llm_server", cmd=server_cmd, @@ -275,17 +280,29 @@ def main_local(): gpu=alloc_mode.gen_pp_size * alloc_mode.gen_tp_size, ) logger.info( - f"LLM inference server launched at: AREAL_LLM_SERVER_ADDRS={','.join(server_addrs)}" + f"SGLang inference servers launched at: AREAL_LLM_SERVER_ADDRS={','.join(server_addrs)}" ) + logger.info("Waiting for SGLang servers to be ready before starting training...") + + # Give servers some time to start up before launching training + # The RemoteSGLangEngine will also do health checks, but this provides clearer logging + time.sleep(2) + logger.info("SGLang servers startup initiated. Training script will verify server readiness.") # Launch trainer entrypoint if not cfg.server_only: + logger.info("=" * 60) + logger.info("STARTUP SEQUENCE: Now starting training script") + logger.info("=" * 60) + logger.info("Training script will connect to the SGLang servers that were just started") + launcher.submit( job_name="trainer", cmd=f"torchrun --nnodes 1 --nproc-per-node {alloc_mode.train_world_size} --master-addr localhost --master-port {find_free_ports(1, (10000, 50000))[0]} {' '.join(sys.argv[1:])}", gpu=alloc_mode.train_world_size, env_vars=dict(AREAL_LLM_SERVER_ADDRS=",".join(server_addrs)), ) + logger.info("Training script launched with SGLang server addresses in environment") try: launcher.wait( diff --git a/docs/arealite/gsm8k_grpo.md b/docs/arealite/gsm8k_grpo.md index 26125d32ed..e1f2862dc2 100644 --- a/docs/arealite/gsm8k_grpo.md +++ b/docs/arealite/gsm8k_grpo.md @@ -6,6 +6,31 @@ the training script configuration file [examples/arealite/configs/gsm8k_grpo.yaml](../../examples/arealite/configs/gsm8k_grpo.yaml). +## SGLang Server Startup Timing + +**Key Question**: When does the SGLang server get started if GRPO training process is triggered by `examples/arealite/gsm8k_grpo.py`? + +**Answer**: SGLang servers are started **BEFORE** the GRPO training script begins execution. + +### Startup Sequence: + +1. **Launcher Invocation**: + ```bash + python -m arealite.launcher.local examples/arealite/gsm8k_grpo.py --config examples/arealite/configs/gsm8k_grpo.yaml + ``` + +2. **SGLang Servers Start First**: The launcher starts SGLang inference servers based on the `allocation_mode` configuration + +3. **Server Readiness**: Launcher waits for all SGLang servers to be healthy and ready + +4. **Environment Setup**: Server addresses are set in `AREAL_LLM_SERVER_ADDRS` environment variable + +5. **Training Script Launch**: Only then is `gsm8k_grpo.py` started with the server addresses + +6. **Connection**: The training script connects to the pre-existing SGLang servers via `RemoteSGLangEngine` + +This design enables AReaL's asynchronous RL training where inference and training are decoupled. + ## How AReaLite Works The following figure illustrates the launching and one asynchronous training step of the @@ -41,6 +66,26 @@ python -m arealite.launcher.ray --config python -m arealite.launcher.slurm --config ``` +### SGLang Server Startup Sequence + +**IMPORTANT**: Understanding the startup timing is crucial for GRPO training: + +1. **SGLang Servers Start FIRST**: The launcher starts SGLang inference servers before + the training script begins execution +2. **Server Readiness Check**: The launcher waits for all SGLang servers to be ready + and responding to health checks +3. **Environment Setup**: Server addresses are collected and passed to the training + script via the `AREAL_LLM_SERVER_ADDRS` environment variable +4. **Training Script Launch**: Only after servers are ready, the training script + (e.g., `gsm8k_grpo.py`) is started +5. **Connection**: The training script connects to the pre-existing SGLang servers + via `RemoteSGLangEngine` + +This sequence ensures that when the GRPO training process begins, the SGLang servers +are already running and ready to handle inference requests. + +### Architecture Details + In AReaLite: - The **training script** is an SPMD python script that serves as the experiment entry diff --git a/examples/arealite/gsm8k_grpo.py b/examples/arealite/gsm8k_grpo.py index f0cfb9e03a..d6d1f5a218 100644 --- a/examples/arealite/gsm8k_grpo.py +++ b/examples/arealite/gsm8k_grpo.py @@ -28,6 +28,20 @@ def gsm8k_reward_fn(prompt, completions, prompt_ids, completion_ids, answer, **k def main(args): + """ + Main GRPO training function for GSM8K dataset. + + Important: This training script expects SGLang inference servers to be + ALREADY RUNNING before this script starts. The servers are launched by + the AReaL launcher (e.g., arealite.launcher.local) which: + + 1. First starts SGLang servers based on allocation_mode config + 2. Waits for servers to be ready and healthy + 3. Sets AREAL_LLM_SERVER_ADDRS environment variable + 4. Then starts this training script + + This script connects to the pre-existing servers via RemoteSGLangEngine. + """ config, _ = load_expr_config(args, GRPOConfig) config: GRPOConfig @@ -77,12 +91,25 @@ def main(args): ) # Initialize inference engine + # NOTE: SGLang servers must already be running at this point. + # They are started by the launcher before this training script begins. + logger.info("Connecting to pre-existing SGLang inference servers...") + server_addrs = os.getenv("AREAL_LLM_SERVER_ADDRS", "").split(",") + if not server_addrs or server_addrs == [""]: + raise RuntimeError( + "No SGLang server addresses found in AREAL_LLM_SERVER_ADDRS. " + "SGLang servers must be started before this training script runs. " + "Use the AReaL launcher (e.g., python -m arealite.launcher.local) to start servers first." + ) + logger.info(f"Found SGLang servers at: {server_addrs}") + rollout = RemoteSGLangEngine(config.rollout) rollout.initialize(None, ft_spec) eval_rollout = RemoteSGLangEngine(config.rollout) eval_rollout.initialize(None, ft_spec) # NOTE: set a large version such that eval does not have any offpolicyness control eval_rollout.set_version(int(1e12)) + logger.info("Successfully connected to SGLang inference servers") # Initialize train engine actor = FSDPPPOActor(config=config.actor)