diff --git a/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh b/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh index 510f7e3bdee..e41cb37767b 100644 --- a/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh +++ b/benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh @@ -27,11 +27,32 @@ MODEL_NAME="${MODEL_NAME:-$MODEL}" CONTAINER=$CONTAINER MOUNTS=$MOUNTS VLLM_CONFIG=$VLLM_CONFIG +SBATCH_TIME="${SBATCH_TIME:-${SBATCH_TIMELIMIT:-04:00:00}}" +# Independent mode starts one complete TP model replica per node. Coupled mode +# forms one multi-node DP/EP engine per tier for models that cannot fit per node. +VLLM_PD_DEPLOYMENT_MODE="${VLLM_PD_DEPLOYMENT_MODE:-independent}" +# Empty falls back to main's SEGMENT or the calculated node count below. +# Coupled deployments can override this with their tier size. +VLLM_SLURM_SEGMENT="${VLLM_SLURM_SEGMENT:-}" SLURM_COMMENT="${SLURM_COMMENT:-}" OPENSANDBOX_DOMAIN="${OPENSANDBOX_DOMAIN:-}" OPENSANDBOX_API_KEY="${OPENSANDBOX_API_KEY:-}" OPENSANDBOX_PROTOCOL="${OPENSANDBOX_PROTOCOL:-http}" +case "$VLLM_PD_DEPLOYMENT_MODE" in + independent | coupled) + ;; + *) + echo "ERROR: VLLM_PD_DEPLOYMENT_MODE must be independent or coupled; got '$VLLM_PD_DEPLOYMENT_MODE'." >&2 + exit 1 + ;; +esac + +if [[ "$VLLM_MODE" == aggregated && "$VLLM_PD_DEPLOYMENT_MODE" == coupled ]]; then + echo "ERROR: VLLM_MODE=aggregated does not support VLLM_PD_DEPLOYMENT_MODE=coupled." >&2 + exit 1 +fi + should_run_eval=$(( $# > 0 )) if (( should_run_eval )); then EXPERIMENT_NAME=$EXPERIMENT_NAME @@ -51,6 +72,11 @@ DECODE_VLLM_NIXL_SIDE_CHANNEL_PORT=5700 ROUTER_SERVER_PORT=8000 WORKER_SERVER_PORT=8001 +PREFILL_SERVER_PORT=8001 +DECODE_SERVER_PORT=8002 + +PREFILL_DP_RPC_PORT=13345 +DECODE_DP_RPC_PORT=13346 ROUTER_PREFILL_POLICY="${ROUTER_PREFILL_POLICY:-cache_aware}" ROUTER_DECODE_POLICY="${ROUTER_DECODE_POLICY:-cache_aware}" @@ -67,6 +93,7 @@ cd /opt/Gym export NEMO_GYM_RUN_ID="\$SLURM_JOB_ID" export NEMO_GYM_USER="\${NEMO_GYM_USER:-\$SLURM_JOB_USER}" +GYM_MODEL_PARAMS=() source "$VLLM_CONFIG" gym eval prepare $@ +use_cached_prepared_benchmarks=true @@ -156,107 +183,237 @@ fi this_node_hostname=\$(hostname) read -r -a nodes <<< "\$ALL_NODES" -router_pid="" -cleanup_vllm() { - if [[ -n "\$router_pid" ]]; then - kill "\$router_pid" 2>/dev/null || true - wait "\$router_pid" 2>/dev/null || true - fi -} -trap cleanup_vllm EXIT -trap 'exit 130' INT -trap 'exit 143' TERM - -if (( SLURM_PROCID == 0 )); then - - # Set a super long request timeout since some reasoning requests may take a long time to generate. - # Don't manually wait as vllm-router will wait for the URLs to come up - # Set a longer worker startup timeout since some models e.g. DSv4 take > 10 mins to load. - router_args=( \ - --host \$this_node_hostname \ - --port $ROUTER_SERVER_PORT \ - --intra-node-data-parallel-size $ROUTER_INTRA_NODE_DATA_PARALLEL_SIZE \ - --request-timeout-secs 86400 \ - --worker-startup-timeout-secs 1200 \ - --log-level error - ) - - if [[ "$VLLM_MODE" == pd ]]; then - router_args+=( \ +if [[ "$VLLM_MODE" == pd && "$VLLM_PD_DEPLOYMENT_MODE" == coupled ]]; then + PREFILL_HEAD=\${nodes[0]} + DECODE_HEAD=\${nodes[$NUM_PREFILL_NODES]} + + wait_for_vllm_health() { + local role=\$1 + local url=\$2 + local local_pid=\${3:-} + local local_role=\${4:-\$role} + + while true; do + if [[ -n "\$local_pid" ]] && ! kill -0 "\$local_pid" 2>/dev/null; then + local status=0 + wait "\$local_pid" || status=\$? + (( status != 0 )) || status=1 + echo "ERROR: \$local_role vLLM process exited while waiting for \$role health (status=\$status)." >&2 + return "\$status" + fi + # Bound each probe so a stalled endpoint cannot block process checks. + # Timeouts retry below; they do not limit overall model startup time. + if curl -fs --connect-timeout 5 --max-time 10 "\$url" >/dev/null; then + return 0 + fi + sleep 5 + done + } + + if (( SLURM_PROCID == 0 )); then + # The first prefill rank owns its tier's API server. The remaining + # prefill ranks run headless so expert parallelism spans the tier. + VLLM_NIXL_SIDE_CHANNEL_HOST=\$this_node_hostname \ + VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_VLLM_NIXL_SIDE_CHANNEL_PORT \ + vllm serve "$MODEL" --served-model-name "$MODEL_NAME" "\${VLLM_COMMON_ARGS[@]}" "\${VLLM_PREFILL_ARGS[@]}" \ + --host \$this_node_hostname \ + --port $PREFILL_SERVER_PORT \ + --data-parallel-size $NUM_PREFILL_NODES \ + --data-parallel-address \$PREFILL_HEAD \ + --data-parallel-rpc-port $PREFILL_DP_RPC_PORT \ + --api-server-count 1 \ + & + prefill_pid=\$! + coupled_pids=("\$prefill_pid") + cleanup_coupled_head() { + local status=\$? + trap - EXIT INT TERM + # Signal both local services without delaying failure propagation; + # the enclosing srun tears down the remaining distributed workers. + kill "\${coupled_pids[@]}" 2>/dev/null || true + exit "\$status" + } + trap cleanup_coupled_head EXIT + trap 'exit 130' INT + trap 'exit 143' TERM + + wait_for_vllm_health "prefill" "http://\$PREFILL_HEAD:$PREFILL_SERVER_PORT/health" "\$prefill_pid" + # Monitor the local prefill process while the remote decode API + # starts. The enclosing srun handles failures on the decode ranks. + wait_for_vllm_health "decode" "http://\$DECODE_HEAD:$DECODE_SERVER_PORT/health" "\$prefill_pid" "prefill" + + vllm-router \ --prefill-policy $ROUTER_PREFILL_POLICY \ --decode-policy $ROUTER_DECODE_POLICY \ - --vllm-pd-disaggregation - ) - for (( i = 0; i < $NUM_PREFILL_NODES; i++ )); do - router_args+=(--prefill "http://\${nodes[i]}:$WORKER_SERVER_PORT") - done - for (( i = 0; i < $NUM_DECODE_NODES; i++ )); do - node_idx=\$(( $NUM_PREFILL_NODES + i )) - router_args+=(--decode "http://\${nodes[node_idx]}:$WORKER_SERVER_PORT") + --vllm-pd-disaggregation \ + --prefill "http://\$PREFILL_HEAD:$PREFILL_SERVER_PORT" \ + --decode "http://\$DECODE_HEAD:$DECODE_SERVER_PORT" \ + --host \$PREFILL_HEAD \ + --port $ROUTER_SERVER_PORT \ + --intra-node-data-parallel-size $ROUTER_INTRA_NODE_DATA_PARALLEL_SIZE \ + --request-timeout-secs 86400 \ + --log-level error & + router_pid=\$! + coupled_pids+=("\$router_pid") + + # Monitor both services after readiness. Polling also catches children that + # exited before monitoring started, which wait -n can otherwise miss. + while kill -0 "\$prefill_pid" 2>/dev/null && kill -0 "\$router_pid" 2>/dev/null; do + sleep 1 done + failed_role=prefill + failed_pid=\$prefill_pid + if kill -0 "\$prefill_pid" 2>/dev/null; then + failed_role=router + failed_pid=\$router_pid + fi + failed_status=0 + wait "\$failed_pid" || failed_status=\$? + # Neither service should exit by itself, even with a zero exit status. + (( failed_status != 0 )) || failed_status=1 + echo "ERROR: \$failed_role process exited after startup (status=\$failed_status)." >&2 + exit "\$failed_status" + elif (( SLURM_PROCID < $NUM_PREFILL_NODES )); then + VLLM_NIXL_SIDE_CHANNEL_HOST=\$this_node_hostname \ + VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_VLLM_NIXL_SIDE_CHANNEL_PORT \ + vllm serve "$MODEL" --served-model-name "$MODEL_NAME" "\${VLLM_COMMON_ARGS[@]}" "\${VLLM_PREFILL_ARGS[@]}" \ + --headless \ + --data-parallel-size $NUM_PREFILL_NODES \ + --data-parallel-start-rank \$SLURM_PROCID \ + --data-parallel-address \$PREFILL_HEAD \ + --data-parallel-rpc-port $PREFILL_DP_RPC_PORT + elif (( SLURM_PROCID == $NUM_PREFILL_NODES )); then + # Decode mirrors prefill with one API rank and headless ranks across + # the other decode nodes. + VLLM_NIXL_SIDE_CHANNEL_HOST=\$this_node_hostname \ + VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_VLLM_NIXL_SIDE_CHANNEL_PORT \ + vllm serve "$MODEL" --served-model-name "$MODEL_NAME" "\${VLLM_COMMON_ARGS[@]}" "\${VLLM_DECODE_ARGS[@]}" \ + --host \$this_node_hostname \ + --port $DECODE_SERVER_PORT \ + --data-parallel-size $NUM_DECODE_NODES \ + --data-parallel-address \$DECODE_HEAD \ + --data-parallel-rpc-port $DECODE_DP_RPC_PORT \ + --api-server-count 1 else - router_args+=(--policy $ROUTER_POLICY --worker-urls) - for node in "\${nodes[@]}"; do - router_args+=("http://\$node:$WORKER_SERVER_PORT") - done + VLLM_NIXL_SIDE_CHANNEL_HOST=\$this_node_hostname \ + VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_VLLM_NIXL_SIDE_CHANNEL_PORT \ + vllm serve "$MODEL" --served-model-name "$MODEL_NAME" "\${VLLM_COMMON_ARGS[@]}" "\${VLLM_DECODE_ARGS[@]}" \ + --headless \ + --data-parallel-size $NUM_DECODE_NODES \ + --data-parallel-start-rank \$(( SLURM_PROCID - $NUM_PREFILL_NODES )) \ + --data-parallel-address \$DECODE_HEAD \ + --data-parallel-rpc-port $DECODE_DP_RPC_PORT fi +else + router_pid="" - vllm-router "\${router_args[@]}" & + cleanup_vllm() { + if [[ -n "\$router_pid" ]]; then + kill "\$router_pid" 2>/dev/null || true + wait "\$router_pid" 2>/dev/null || true + fi + } + trap cleanup_vllm EXIT + trap 'exit 130' INT + trap 'exit 143' TERM + + if (( SLURM_PROCID == 0 )); then + + # Set a super long request timeout since some reasoning requests may take a long time to generate. + # Don't manually wait as vllm-router will wait for the URLs to come up + # Set a longer worker startup timeout since some models e.g. DSv4 take > 10 mins to load. + router_args=( \ + --host \$this_node_hostname \ + --port $ROUTER_SERVER_PORT \ + --intra-node-data-parallel-size $ROUTER_INTRA_NODE_DATA_PARALLEL_SIZE \ + --request-timeout-secs 86400 \ + --worker-startup-timeout-secs 1200 \ + --log-level error + ) - router_pid=\$! + if [[ "$VLLM_MODE" == pd ]]; then + router_args+=( \ + --prefill-policy $ROUTER_PREFILL_POLICY \ + --decode-policy $ROUTER_DECODE_POLICY \ + --vllm-pd-disaggregation + ) + for (( i = 0; i < $NUM_PREFILL_NODES; i++ )); do + router_args+=(--prefill "http://\${nodes[i]}:$WORKER_SERVER_PORT") + done + for (( i = 0; i < $NUM_DECODE_NODES; i++ )); do + node_idx=\$(( $NUM_PREFILL_NODES + i )) + router_args+=(--decode "http://\${nodes[node_idx]}:$WORKER_SERVER_PORT") + done + else + router_args+=(--policy $ROUTER_POLICY --worker-urls) + for node in "\${nodes[@]}"; do + router_args+=("http://\$node:$WORKER_SERVER_PORT") + done + fi - sleep 5 - if ! kill -0 "\$router_pid" 2>/dev/null; then - echo "vllm-router exited during startup" >&2 - exit 1 + vllm-router "\${router_args[@]}" & + + router_pid=\$! + + sleep 5 + if ! kill -0 "\$router_pid" 2>/dev/null; then + echo "vllm-router exited during startup" >&2 + exit 1 + fi fi -fi -if [[ "$VLLM_MODE" == aggregated ]]; then - # Reuse decode tuning by default. Configs can supply VLLM_AGGREGATED_ARGS - # instead (including an empty array) for independent tuning. - if declare -p VLLM_AGGREGATED_ARGS &>/dev/null; then - worker_args=("\${VLLM_COMMON_ARGS[@]}" "\${VLLM_AGGREGATED_ARGS[@]}") + if [[ "$VLLM_MODE" == aggregated ]]; then + # Reuse prefill tuning by default. Configs can supply VLLM_AGGREGATED_ARGS + # instead (including an empty array) for independent tuning. + if declare -p VLLM_AGGREGATED_ARGS &>/dev/null; then + worker_args=("\${VLLM_COMMON_ARGS[@]}" "\${VLLM_AGGREGATED_ARGS[@]}") + else + worker_args=("\${VLLM_COMMON_ARGS[@]}" "\${VLLM_PREFILL_ARGS[@]}") + fi + # A complete replica must not wait for KV transfers from a prefill worker. + # Strip both CLI forms, including connectors set in VLLM_COMMON_ARGS. + aggregated_args=() + skip_value=0 + for arg in "\${worker_args[@]}"; do + if (( skip_value )); then + skip_value=0 + continue + fi + if [[ "\$arg" == --kv-transfer-config ]]; then + skip_value=1 + elif [[ "\$arg" != --kv-transfer-config=* ]]; then + aggregated_args+=("\$arg") + fi + done + vllm serve "$MODEL" --served-model-name "$MODEL_NAME" "\${aggregated_args[@]}" \ + --host \$this_node_hostname \ + --port $WORKER_SERVER_PORT + elif (( SLURM_PROCID < $NUM_PREFILL_NODES )); then + # Prefill + VLLM_NIXL_SIDE_CHANNEL_HOST=\$this_node_hostname \ + VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_VLLM_NIXL_SIDE_CHANNEL_PORT \ + vllm serve "$MODEL" --served-model-name "$MODEL_NAME" "\${VLLM_COMMON_ARGS[@]}" "\${VLLM_PREFILL_ARGS[@]}" \ + --host \$this_node_hostname \ + --port $WORKER_SERVER_PORT else - worker_args=("\${VLLM_COMMON_ARGS[@]}" "\${VLLM_PREFILL_ARGS[@]}") + # Decode + VLLM_NIXL_SIDE_CHANNEL_HOST=\$this_node_hostname \ + VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_VLLM_NIXL_SIDE_CHANNEL_PORT \ + vllm serve "$MODEL" --served-model-name "$MODEL_NAME" "\${VLLM_COMMON_ARGS[@]}" "\${VLLM_DECODE_ARGS[@]}" \ + --host \$this_node_hostname \ + --port $WORKER_SERVER_PORT fi - # A complete replica must not wait for KV transfers from a prefill worker. - # Strip both CLI forms, including connectors set in VLLM_COMMON_ARGS. - aggregated_args=() - skip_value=0 - for arg in "\${worker_args[@]}"; do - if (( skip_value )); then - skip_value=0 - continue - fi - if [[ "\$arg" == --kv-transfer-config ]]; then - skip_value=1 - elif [[ "\$arg" != --kv-transfer-config=* ]]; then - aggregated_args+=("\$arg") - fi - done - vllm serve "$MODEL" --served-model-name "$MODEL_NAME" "\${aggregated_args[@]}" \ - --host \$this_node_hostname \ - --port $WORKER_SERVER_PORT -elif (( SLURM_PROCID < $NUM_PREFILL_NODES )); then - # Prefill - VLLM_NIXL_SIDE_CHANNEL_HOST=\$this_node_hostname \ - VLLM_NIXL_SIDE_CHANNEL_PORT=$PREFILL_VLLM_NIXL_SIDE_CHANNEL_PORT \ - vllm serve "$MODEL" --served-model-name "$MODEL_NAME" "\${VLLM_COMMON_ARGS[@]}" "\${VLLM_PREFILL_ARGS[@]}" \ - --host \$this_node_hostname \ - --port $WORKER_SERVER_PORT -else - # Decode - VLLM_NIXL_SIDE_CHANNEL_HOST=\$this_node_hostname \ - VLLM_NIXL_SIDE_CHANNEL_PORT=$DECODE_VLLM_NIXL_SIDE_CHANNEL_PORT \ - vllm serve "$MODEL" --served-model-name "$MODEL_NAME" "\${VLLM_COMMON_ARGS[@]}" "\${VLLM_DECODE_ARGS[@]}" \ - --host \$this_node_hostname \ - --port $WORKER_SERVER_PORT fi EOF ) +VLLM_SLURM_SEGMENT="${VLLM_SLURM_SEGMENT:-${SEGMENT:-$NUM_NODES}}" +if [[ ! "$VLLM_SLURM_SEGMENT" =~ ^[1-9][0-9]*$ ]]; then + echo "ERROR: VLLM_SLURM_SEGMENT must be a positive integer." >&2 + exit 2 +fi + batch_command=$(cat < 0 otherwise the engine will hang on the second or third engine step. -SEGMENT=${SEGMENT:-$NUM_NODES} - +# This cluster needs --segment > 0 to avoid distributed engine hangs. +# Coupled tiers benefit from caller-configurable segments matching their node count. submit_dir=$(pwd -P) # An exported connection is sent as arguments; otherwise env.yaml is read. if [[ -n "$OPENSANDBOX_DOMAIN" ]]; then @@ -354,13 +510,13 @@ main_job_id=$( sbatch \ --parsable \ --nodes=$NUM_NODES \ - --time="${SBATCH_TIMELIMIT:-04:00:00}" \ + --time="$SBATCH_TIME" \ + --segment="$VLLM_SLURM_SEGMENT" \ --job-name=gym-$EXPERIMENT_NAME-$USER \ --output=slurm-logs/%j-%x.log \ --ntasks-per-node=1 \ --comment="$SLURM_COMMENT" \ --exclusive \ - --segment=$SEGMENT \ --wrap 'exec bash -c "$batch_command"' ) main_job_id=${main_job_id%%;*} diff --git a/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh b/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh new file mode 100644 index 00000000000..cd49feba330 --- /dev/null +++ b/benchmarks/nemotron_3.5_super/vllm_configs/nemotron_3_ultra.sh @@ -0,0 +1,78 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Nemotron 3 Ultra BF16 configuration for disaggregated prefill/decode on +# 4-GPU GB200 nodes. Each tier uses four coupled data-parallel ranks so expert +# parallelism can shard the 512 experts over 16 GPUs. Launch this config with +# VLLM_PD_DEPLOYMENT_MODE=coupled and VLLM_SLURM_SEGMENT=4. + +# Both tiers use piecewise CUDA graphs with graph-owned inputs and MTP3. +# Settings are fixed in this recipe; asynchronous scheduling is prefill-only. + +# Mamba state transfers require the dimension-sequence layout. +export VLLM_SSM_CONV_STATE_LAYOUT=DS +# The V1 model runner provides higher decode throughput for Ultra. +export VLLM_USE_V2_MODEL_RUNNER=0 + +# Ultra uses the GB200 InfiniBand interface and NCCL's default configuration. +# Recipe settings override the shared launcher's communication defaults. +export UCX_TLS=rc_x,rc,cuda_copy,cuda_ipc +export UCX_NET_DEVICES=mlx5_0:1 +export UCX_IB_ADDR_TYPE=eth +unset NCCL_CUMEM_ENABLE NCCL_MNNVL_ENABLE NCCL_NVLS_ENABLE + +# Standard safetensors loading avoided the InstantTensor io_uring failures seen +# against the Lustre-hosted checkpoint. +export SAFETENSORS_FAST_GPU=1 + +VLLM_COMMON_ARGS=( + --disable-uvicorn-access-log + --trust-remote-code + --dtype bfloat16 + --distributed-executor-backend mp + --data-parallel-backend mp + --max-model-len 262144 + --enable-auto-tool-choice + --tool-call-parser qwen3_coder + --reasoning-parser nemotron_v3 + --enable-chunked-prefill + --kv-cache-dtype fp8 + --no-disable-hybrid-kv-cache-manager + --block-size 128 + --mamba-cache-mode align + --mamba-ssm-cache-dtype float16 + --mamba-backend flashinfer + --enable-mamba-cache-stochastic-rounding + --mamba-cache-philox-rounds 5 + --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 96}' + --load-format safetensors + --enable-expert-parallel + --distributed-timeout-seconds 3600 + --enable-prefix-caching + # Both tiers need the same speculative width for compatible cache layouts. + --speculative-config '{"method":"mtp","num_speculative_tokens":3}' +) + +VLLM_PREFILL_ARGS=( + --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_producer","kv_load_failure_policy":"fail"}' + --gpu-memory-utilization 0.90 + --max-num-batched-tokens 16384 + --max-num-seqs 64 + --data-parallel-size-local 1 + --tensor-parallel-size 4 + --async-scheduling + --compilation-config '{"cudagraph_mode":"PIECEWISE","cudagraph_copy_inputs":true,"pass_config":{"fuse_allreduce_rms":false}}' +) + +VLLM_DECODE_ARGS=( + --kv-transfer-config '{"kv_connector":"NixlConnector","kv_role":"kv_consumer","kv_load_failure_policy":"fail"}' + --compilation-config '{"cudagraph_mode":"PIECEWISE","cudagraph_copy_inputs":true,"cudagraph_capture_sizes": + [1,2,3,4,5,8,10,12,15,16,20,24,25,28,30,32,35,40,45,50,55,60,65,70,75,80,128,256,512],"pass_config":{"fuse_allreduce_rms":false}}' + --gpu-memory-utilization 0.95 + --max-num-batched-tokens 8192 + --max-num-seqs 64 + --data-parallel-size-local 1 + --tensor-parallel-size 4 + --no-async-scheduling +) diff --git a/tests/unit_tests/test_super_vllm_launcher.py b/tests/unit_tests/test_super_vllm_launcher.py new file mode 100644 index 00000000000..99830440967 --- /dev/null +++ b/tests/unit_tests/test_super_vllm_launcher.py @@ -0,0 +1,774 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os +import signal +import subprocess +import unittest +from pathlib import Path +from tempfile import TemporaryDirectory + + +SCRIPT = Path(__file__).resolve().parents[2] / "benchmarks/nemotron_3.5_super/sbatch_external_vllm.sh" + + +class TestSuperVllmLauncher(unittest.TestCase): + def setUp(self): + # Never inherit cluster credentials, tuning overrides, or real sbatch commands. + self.env = { + "PATH": os.defpath, + "USER": "launcher-test", + "MODEL": "/test/model", + "CONTAINER": "/test/image.sqsh", + "MOUNTS": "/test:/test", + "VLLM_CONFIG": "/dev/null", + "EXPERIMENT_NAME": "launcher-test", + "NUM_PREFILL_NODES": "4", + "NUM_DECODE_NODES": "4", + "SLURM_PROCID": "0", + "SLURM_JOB_ID": "12345", + "SLURM_JOB_USER": "launcher-test", + "ROUTER_NODE": "node0", + "ALL_NODES": "node0 node1 node2 node3 node4 node5 node6 node7", + } + + def run_shell(self, command, *args, env=None): + proc = subprocess.Popen( + ["bash", "--noprofile", "--norc", "-c", command, "launcher-test", *args], + env=self.env | (env or {}), + stdin=subprocess.DEVNULL, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + start_new_session=True, + ) + try: + stdout, stderr = proc.communicate(timeout=5) + except subprocess.TimeoutExpired: + # The failure-path tests must terminate instead of leaking a polling worker. + os.killpg(proc.pid, signal.SIGKILL) + stdout, stderr = proc.communicate() + self.fail(f"Launcher did not terminate. stdout={stdout!r}, stderr={stderr!r}") + return proc.returncode, stdout, stderr + + def capture_submission(self, *eval_args, env=None): + # Capture generated commands and both sbatch calls without submitting jobs. + stub = r""" +sbatch() { + if [[ -n "${vllm_command:-}" ]]; then + printf '%s\0' "$eval_command" "$vllm_command" "$batch_command" >&2 + fi + printf '%s\0' "$@" >&2 + printf '\0' >&2 + printf '12345\n' +} +launcher_script=$1 +shift +source "$launcher_script" "$@" +""" + status, _, captured = self.run_shell(stub, str(SCRIPT), *eval_args, env=env) + self.assertEqual(status, 0, captured) + eval_command, pd_command, batch_command, submissions = captured.split("\0", 3) + self.assertTrue(submissions.endswith("\0\0")) + calls = [call.split("\0") for call in submissions.removesuffix("\0\0").split("\0\0")] + return eval_command, pd_command, batch_command, calls + + def generate_commands(self, *overrides, env=None): + eval_command, pd_command, _, _ = self.capture_submission("--config", "benchmark.yaml", *overrides, env=env) + return eval_command, pd_command + + def eval_arguments(self, *overrides, env=None): + command, _ = self.generate_commands(*overrides, env=env) + command = command.replace("source /opt/Gym_venv/bin/activate", ":").replace("cd /opt/Gym\n", ":\n") + stubs = r""" +gym() { + if [[ "$2" == run ]]; then printf '%s\0' "$@"; fi +} +date() { printf '%s\n' "${TEST_DATE:-20260909_120000}"; } +getent() { printf '10.0.0.1 node0\n'; } +""" + status, stdout, stderr = self.run_shell(stubs + command, env=env) + self.assertEqual(status, 0, stderr) + return stdout.rstrip("\0").split("\0") + + def settings(self, args, key): + return [arg for arg in args if arg.lstrip("+").startswith(key + "=")] + + def serving_arguments(self, command, *, rank, coupled_head=False, env=None): + # Record argv separately for vLLM and the router; marker files synchronize startup. + stubs = r""" +VLLM_COMMON_ARGS=(--common-test 'value with spaces') +VLLM_PREFILL_ARGS=(--prefill-test producer) +VLLM_DECODE_ARGS=(--decode-test consumer) +vllm() { + printf '%s\0' "$VLLM_NIXL_SIDE_CHANNEL_HOST" "$VLLM_NIXL_SIDE_CHANNEL_PORT" "$@" + touch "$TEST_STATE_DIR/service-ready" + if [[ "$TEST_COUPLED_HEAD" == 1 ]]; then + while true; do command sleep 0.01; done + fi + if (( SLURM_PROCID == 0 )); then + while [[ ! -f "$TEST_STATE_DIR/router-ready" ]]; do command sleep 0.01; done + fi +} +vllm-router() { + printf '%s\0' "$@" >&2 + touch "$TEST_STATE_DIR/router-ready" + if [[ "$TEST_COUPLED_HEAD" == 1 ]]; then + kill -TERM "$$" + else + # Independent mode checks router liveness before starting vLLM. + while true; do command sleep 0.01; done + fi +} +curl() { [[ -f "$TEST_STATE_DIR/service-ready" ]]; } +sleep() { command sleep 0.01; } +hostname() { printf 'node%s\n' "$SLURM_PROCID"; } +""" + with TemporaryDirectory(prefix="gym-serving-args-") as state_dir: + status, stdout, stderr = self.run_shell( + stubs + command, + env=(env or {}) + | { + "SLURM_PROCID": str(rank), + "TEST_STATE_DIR": state_dir, + "TEST_COUPLED_HEAD": str(int(coupled_head)), + }, + ) + self.assertEqual(status, 143 if coupled_head else 0, stderr) + host, nixl_port, *vllm_args = stdout.rstrip("\0").split("\0") + router_args = stderr.rstrip("\0").split("\0") if stderr else [] + return host, nixl_port, vllm_args, router_args + + def assert_router_arguments( + self, args: list[str], prefill_urls: list[str], decode_urls: list[str], *, startup_timeout: bool = False + ) -> None: + expected = { + "--prefill-policy": "cache_aware", + "--decode-policy": "cache_aware", + "--host": "node0", + "--port": "8000", + "--intra-node-data-parallel-size": "1", + "--request-timeout-secs": "86400", + "--log-level": "error", + } + if startup_timeout: + expected["--worker-startup-timeout-secs"] = "1200" + self.assertEqual(args.count("--vllm-pd-disaggregation"), 1) + args = [arg for arg in args if arg != "--vllm-pd-disaggregation"] + # URL options may precede or follow the common options; retain their tier order. + for flag, urls in (("--prefill", prefill_urls), ("--decode", decode_urls)): + actual_urls = [] + remaining = [] + i = 0 + while i < len(args): + if args[i] == flag: + actual_urls.append(args[i + 1]) + i += 2 + else: + remaining.append(args[i]) + i += 1 + self.assertEqual(actual_urls, urls) + args = remaining + self.assertEqual(len(args), 2 * len(expected)) + self.assertEqual(dict(zip(args[::2], args[1::2], strict=True)), expected) + + def test_independent_mode_preserves_per_node_engines(self) -> None: + """Independent mode runs one engine per node and routes requests to each tier.""" + for mode in (None, "independent"): + for prefill_count, decode_count in ((1, 1), (1, 4), (4, 4)): + env = {"NUM_PREFILL_NODES": str(prefill_count), "NUM_DECODE_NODES": str(decode_count)} + if mode is not None: + env["VLLM_PD_DEPLOYMENT_MODE"] = mode + _, command = self.generate_commands(env=env) + for rank in range(prefill_count + decode_count): + with self.subTest(mode=mode, prefill=prefill_count, decode=decode_count, rank=rank): + host, nixl_port, args, router = self.serving_arguments(command, rank=rank, env=env) + is_prefill = rank < prefill_count + self.assertEqual(host, f"node{rank}") + self.assertEqual(nixl_port, "5600" if is_prefill else "5700") + self.assertEqual( + args, + [ + "serve", + "/test/model", + "--served-model-name", + "/test/model", + "--common-test", + "value with spaces", + "--prefill-test" if is_prefill else "--decode-test", + "producer" if is_prefill else "consumer", + "--host", + f"node{rank}", + "--port", + "8001", + ], + ) + if rank == 0: + self.assert_router_arguments( + router, + [f"http://node{i}:8001" for i in range(prefill_count)], + [f"http://node{i}:8001" for i in range(prefill_count, prefill_count + decode_count)], + startup_timeout=True, + ) + else: + self.assertEqual(router, []) + + def test_aggregated_submission_preserves_node_count(self) -> None: + """Aggregated allocation uses NUM_NODES, including its one-node default.""" + for count in (None, 3): + for eval_args in ((), ("--config", "benchmark.yaml")): + with self.subTest(count=count, evaluation=bool(eval_args)): + env = {"VLLM_MODE": "aggregated", "NUM_PREFILL_NODES": "", "NUM_DECODE_NODES": ""} + if count is not None: + env["NUM_NODES"] = str(count) + _, _, batch, calls = self.capture_submission(*eval_args, env=env) + self.assertEqual(len(calls), 2 if eval_args else 1) + self.assertIn(f"--nodes={count or 1}", calls[0]) + self.assertIn(f"--segment={count or 1}", calls[0]) + self.assertIn(f"--ntasks={count or 1}", batch) + + def test_aggregated_coupled_combination_is_rejected(self) -> None: + """Reject unsupported multi-node aggregated execution before submitting jobs.""" + status, stdout, stderr = self.run_shell( + 'sbatch() { printf "unexpected-submission\\n"; }; source "$@"', + str(SCRIPT), + env={"VLLM_MODE": "aggregated", "VLLM_PD_DEPLOYMENT_MODE": "coupled"}, + ) + self.assertEqual(status, 1) + self.assertNotIn("unexpected-submission", stdout) + self.assertIn("VLLM_MODE=aggregated does not support VLLM_PD_DEPLOYMENT_MODE=coupled", stderr) + + def test_generated_scripts_have_valid_syntax(self) -> None: + """Parse the generated scripts too: outer bash -n cannot validate heredoc contents.""" + for env in ({}, {"VLLM_PD_DEPLOYMENT_MODE": "coupled"}, {"VLLM_MODE": "aggregated"}): + with self.subTest(env=env): + evaluation, serving, batch, _ = self.capture_submission("--config", "benchmark.yaml", env=env) + for command in (evaluation, serving, batch): + result = subprocess.run(["bash", "-n"], input=command, text=True, capture_output=True, timeout=5) + self.assertEqual(result.returncode, 0, result.stderr) + + def test_ultra_runtime_settings_override_launcher_defaults(self) -> None: + """Ultra configures its model runner, state layout, and communication environment.""" + expected = { + "VLLM_USE_V2_MODEL_RUNNER": "0", + "VLLM_SSM_CONV_STATE_LAYOUT": "DS", + "UCX_TLS": "rc_x,rc,cuda_copy,cuda_ipc", + "UCX_NET_DEVICES": "mlx5_0:1", + "UCX_IB_ADDR_TYPE": "eth", + "NCCL_CUMEM_ENABLE": "unset", + "NCCL_MNNVL_ENABLE": "unset", + "NCCL_NVLS_ENABLE": "unset", + } + recipe = str(SCRIPT.parent / "vllm_configs/nemotron_3_ultra.sh") + _, command = self.generate_commands(env={"VLLM_CONFIG": recipe}) + setup = command.split("this_node_hostname=", 1)[0] + inspect = '\nfor name in "$@"; do printf "%s\\0" "${!name-unset}"; done\n' + status, stdout, stderr = self.run_shell(setup + inspect, *expected) + self.assertEqual(status, 0, stderr) + self.assertEqual(dict(zip(expected, stdout.removesuffix("\0").split("\0"), strict=True)), expected) + + def test_coupled_nodes_use_correct_tier_roles_and_ranks(self): + """Assign coupled tier roles, ranks, and ports while leaving router balancing thresholds at defaults.""" + for prefill_count, decode_count in ((1, 1), (1, 4), (2, 3), (4, 4)): + env = { + "VLLM_PD_DEPLOYMENT_MODE": "coupled", + "NUM_PREFILL_NODES": str(prefill_count), + "NUM_DECODE_NODES": str(decode_count), + "MODEL_NAME": "served-model-alias", + } + _, command = self.generate_commands(env=env) + for rank in range(prefill_count + decode_count): + with self.subTest(prefill=prefill_count, decode=decode_count, rank=rank): + host, nixl_port, args, router = self.serving_arguments( + command, rank=rank, coupled_head=rank == 0, env=env + ) + is_prefill = rank < prefill_count + local_rank = rank if is_prefill else rank - prefill_count + head = "node0" if is_prefill else f"node{prefill_count}" + self.assertEqual(host, f"node{rank}") + self.assertEqual(nixl_port, "5600" if is_prefill else "5700") + expected = [ + "serve", + "/test/model", + "--served-model-name", + "served-model-alias", + "--common-test", + "value with spaces", + "--prefill-test" if is_prefill else "--decode-test", + "producer" if is_prefill else "consumer", + ] + if local_rank == 0: + expected += ["--host", host, "--port", "8001" if is_prefill else "8002"] + else: + expected += ["--headless"] + expected += ["--data-parallel-size", str(prefill_count if is_prefill else decode_count)] + if local_rank != 0: + expected += ["--data-parallel-start-rank", str(local_rank)] + expected += [ + "--data-parallel-address", + head, + "--data-parallel-rpc-port", + "13345" if is_prefill else "13346", + ] + if local_rank == 0: + expected += ["--api-server-count", "1"] + self.assertEqual(args, expected) + if rank == 0: + self.assert_router_arguments( + router, ["http://node0:8001"], [f"http://node{prefill_count}:8002"] + ) + else: + self.assertEqual(router, []) + + def test_router_overrides_apply_to_both_deployment_modes(self): + """Apply router policy and local data-parallel overrides in both deployment modes.""" + for mode in ("independent", "coupled"): + with self.subTest(mode=mode): + env = { + "VLLM_PD_DEPLOYMENT_MODE": mode, + "ROUTER_PREFILL_POLICY": "round_robin", + "ROUTER_DECODE_POLICY": "random", + "ROUTER_INTRA_NODE_DATA_PARALLEL_SIZE": "2", + } + _, command = self.generate_commands(env=env) + _, _, _, router = self.serving_arguments(command, rank=0, coupled_head=mode == "coupled", env=env) + for flag, value in ( + ("--prefill-policy", "round_robin"), + ("--decode-policy", "random"), + ("--intra-node-data-parallel-size", "2"), + ): + self.assertEqual(router.count(flag), 1) + self.assertEqual(router[router.index(flag) + 1], value) + + def test_independent_router_startup_failure_stops_launch(self): + """Stop independent startup when the router exits during main's five-second startup check.""" + stubs = r""" +VLLM_COMMON_ARGS=() +VLLM_PREFILL_ARGS=() +VLLM_DECODE_ARGS=() +vllm-router() { return "$TEST_ROUTER_STATUS"; } +vllm() { printf 'unexpected-vllm-start\n'; } +hostname() { printf 'node0\n'; } +# Reap the failed router deterministically instead of waiting five real seconds. +sleep() { printf 'startup-delay=%s\n' "$1"; wait "$router_pid" || true; } +""" + _, command = self.generate_commands() + for router_status in (0, 7): + with self.subTest(router_status=router_status): + status, stdout, stderr = self.run_shell( + stubs + command, env={"TEST_ROUTER_STATUS": str(router_status)} + ) + self.assertEqual(status, 1, stderr) + self.assertEqual(stdout, "startup-delay=5\n") + self.assertIn("vllm-router exited during startup", stderr) + + def test_existing_recipes_preserve_sampling_overrides(self): + """Pass through existing models' sampling parameters without injecting Ultra evaluation defaults.""" + for name in ("inkling_small.sh", "qwen3.5-122b-a10b.sh", "nemotron_3.5_super.sh", "nemotron_3.5_lightning.sh"): + with self.subTest(recipe=name): + args = self.eval_arguments(env={"VLLM_CONFIG": str(SCRIPT.parent / "vllm_configs" / name)}) + expected = ( + [] + if name == "inkling_small.sh" + else [ + "++policy_model.responses_api_models.vllm_model.sampling_overrides.temperature=1.0", + "++policy_model.responses_api_models.vllm_model.sampling_overrides.top_p=0.95", + ] + ) + self.assertEqual([arg for arg in args if ".sampling_overrides." in arg], expected) + self.assertEqual(self.settings(args, "num_samples_in_parallel"), []) + self.assertEqual(self.settings(args, "resume_from_cache"), []) + + def test_submission_defaults_preserve_time_and_full_allocation_segment(self): + """Default submissions use a four-hour limit and one segment covering all allocated nodes.""" + for prefill_count, decode_count in ((1, 1), (1, 4), (4, 4)): + with self.subTest(prefill=prefill_count, decode=decode_count): + _, _, _, calls = self.capture_submission( + "--config", + "benchmark.yaml", + env={"NUM_PREFILL_NODES": str(prefill_count), "NUM_DECODE_NODES": str(decode_count)}, + ) + self.assertEqual(len(calls), 2) + self.assertIn(f"--nodes={prefill_count + decode_count}", calls[0]) + self.assertEqual( + [arg for arg in calls[0] if arg.startswith("--segment=")], + [f"--segment={prefill_count + decode_count}"], + ) + self.assertIn("--time=04:00:00", calls[0]) + self.assertIn("--ntasks-per-node=1", calls[0]) + self.assertIn("--exclusive", calls[0]) + self.assertIn("--dependency=afterany:12345", calls[1]) + + def test_segment_controls_resolve_to_one_submission_argument(self): + """Prefer the Ultra segment override, then main's SEGMENT, then calculated nodes, without duplicates.""" + cases = ( + ({}, "8"), + ({"NUM_NODES": "99"}, "8"), + ({"SEGMENT": "2"}, "2"), + ({"VLLM_SLURM_SEGMENT": "4"}, "4"), + ({"VLLM_SLURM_SEGMENT": "4", "SEGMENT": "2"}, "4"), + ({"VLLM_SLURM_SEGMENT": "", "SEGMENT": "2"}, "2"), + ({"VLLM_SLURM_SEGMENT": "4", "SEGMENT": ""}, "4"), + ({"VLLM_SLURM_SEGMENT": "", "SEGMENT": ""}, "8"), + ({"VLLM_SLURM_SEGMENT": "4", "SEGMENT": "invalid"}, "4"), + ) + for mode in ("independent", "coupled"): + for eval_args in ((), ("--config", "benchmark.yaml")): + for segment_env, expected in cases: + with self.subTest(mode=mode, evaluation=bool(eval_args), controls=segment_env): + _, _, _, calls = self.capture_submission( + *eval_args, env={"VLLM_PD_DEPLOYMENT_MODE": mode} | segment_env + ) + self.assertEqual(len(calls), 2 if eval_args else 1) + self.assertEqual( + [arg for arg in calls[0] if arg.startswith("--segment=")], [f"--segment={expected}"] + ) + if eval_args: + self.assertFalse(any(arg.startswith("--segment=") for arg in calls[1])) + + def test_submission_overrides_do_not_change_cleanup_allocation(self): + """Apply custom walltime and segment size only to the main job, leaving cleanup CPU-only and short.""" + _, _, _, calls = self.capture_submission( + "--config", "benchmark.yaml", env={"SBATCH_TIME": "7-00:00:00", "VLLM_SLURM_SEGMENT": "4"} + ) + self.assertEqual(len(calls), 2) + self.assertIn("--time=7-00:00:00", calls[0]) + self.assertEqual([arg for arg in calls[0] if arg.startswith("--segment=")], ["--segment=4"]) + self.assertIn("--time=00:30:00", calls[1]) + self.assertIn("--nodes=1", calls[1]) + self.assertIn("--partition=cpu", calls[1]) + self.assertIn("--qos=cpu-normal", calls[1]) + self.assertIn("--gres=none", calls[1]) + self.assertFalse(any(arg.startswith("--segment=") for arg in calls[1])) + + def test_walltime_alias_precedence_and_cleanup_isolation(self) -> None: + """SBATCH_TIME takes precedence over SBATCH_TIMELIMIT; cleanup has its own walltime.""" + for overrides, expected in ( + ({"SBATCH_TIMELIMIT": "06:00:00"}, "06:00:00"), + ({"SBATCH_TIME": "", "SBATCH_TIMELIMIT": "06:00:00"}, "06:00:00"), + ({"SBATCH_TIME": "7-00:00:00", "SBATCH_TIMELIMIT": "06:00:00"}, "7-00:00:00"), + ({"SBATCH_TIME": "", "SBATCH_TIMELIMIT": ""}, "04:00:00"), + ): + for mode in ("pd", "aggregated"): + with self.subTest(overrides=overrides, mode=mode): + _, _, _, calls = self.capture_submission( + "--config", "benchmark.yaml", env={"VLLM_MODE": mode} | overrides + ) + self.assertEqual([arg for arg in calls[0] if arg.startswith("--time=")], [f"--time={expected}"]) + self.assertIn("--time=00:30:00", calls[1]) + self.assertFalse(any(arg.startswith("--segment=") for arg in calls[1])) + + def test_invalid_launcher_controls_are_rejected_before_submission(self): + """Reject invalid deployment modes and segment sizes without submitting any job.""" + stub = 'sbatch() { printf "unexpected-submission\\n"; }; source "$@"' + cases = { + "VLLM_PD_DEPLOYMENT_MODE": (("bad", "COUPLED"), 1), + "VLLM_SLURM_SEGMENT": (("0", "-1", "1.5", "bad"), 2), + "SEGMENT": (("0", "-1", "1.5", "bad"), 2), + } + for name, (values, expected_status) in cases.items(): + for value in values: + with self.subTest(name=name, value=value): + status, stdout, stderr = self.run_shell(stub, str(SCRIPT), env={name: value}) + self.assertEqual(status, expected_status, stderr) + self.assertNotIn("unexpected-submission", stdout) + self.assertIn(name, stderr) + + def test_invalid_segment_override_does_not_fall_back_to_main(self): + """Reject an invalid explicit Ultra segment even when main's SEGMENT supplies a valid fallback.""" + stub = 'sbatch() { printf "unexpected-submission\\n"; }; source "$@"' + for value in ("0", "-1", "1.5", "bad"): + with self.subTest(value=value): + status, stdout, stderr = self.run_shell( + stub, str(SCRIPT), env={"VLLM_SLURM_SEGMENT": value, "SEGMENT": "4"} + ) + self.assertEqual(status, 2, stderr) + self.assertNotIn("unexpected-submission", stdout) + self.assertIn("VLLM_SLURM_SEGMENT must be a positive integer", stderr) + + def test_serving_only_skips_evaluation_and_cleanup_submission(self): + """Run only the serving step without eval arguments and preserve its success or failure status.""" + stubs = r""" +scontrol() { printf '%s\n' node0 node1 node2 node3 node4 node5 node6 node7; } +srun() { printf '%s\0' "$@"; printf '\0'; return "$TEST_SERVER_STATUS"; } +""" + for mode in ("independent", "coupled"): + for server_status in (0, 7): + with self.subTest(mode=mode, server_status=server_status): + _, _, batch_command, calls = self.capture_submission( + env={"VLLM_PD_DEPLOYMENT_MODE": mode, "EXPERIMENT_NAME": "", "EXPORT_TO_CSV": "1"} + ) + self.assertEqual(len(calls), 1) + self.assertIn("--job-name=gym-vllm_only-launcher-test", calls[0]) + status, stdout, stderr = self.run_shell( + stubs + batch_command, + env={ + "SLURM_JOB_NODELIST": "test-nodes", + "SLURM_SUBMIT_DIR": "/test", + "SLURM_CPUS_ON_NODE": "64", + "vllm_command": "fake-serving-command", + "eval_command": "fake-eval-command", + "TEST_SERVER_STATUS": str(server_status), + }, + ) + self.assertEqual(status, server_status, stderr) + steps = stdout.removesuffix("\0\0").split("\0\0") + self.assertEqual(len(steps), 1) + args = steps[0].split("\0") + self.assertIn("--nodes=8", args) + self.assertIn("--ntasks=8", args) + self.assertIn("--kill-on-bad-exit=1", args) + self.assertIn("fake-serving-command", args) + self.assertNotIn("--overlap", args) + + def test_default_evaluation_settings_are_unchanged(self): + """Leave concurrency and resume to Gym and preserve timestamped output naming by default.""" + args = self.eval_arguments() + self.assertEqual(self.settings(args, "num_samples_in_parallel"), []) + self.assertEqual(self.settings(args, "resume_from_cache"), []) + self.assertIn("++output_jsonl_fpath=results/launcher-test/slurm_job_id_12345/date_20260909_120000.jsonl", args) + + def test_explicit_concurrency_arguments_are_preserved(self): + """Pass explicit concurrency values through to Gym without replacing or duplicating them.""" + for prefix in ("", "+", "++"): + for value in (16, 32, 512): + with self.subTest(prefix=prefix, value=value): + override = f"{prefix}num_samples_in_parallel={value}" + args = self.eval_arguments(override) + self.assertEqual(self.settings(args, "num_samples_in_parallel"), [override]) + + def test_explicit_resume_arguments_are_preserved(self): + """Pass explicit resume values through to Gym without changing the supplied output path.""" + for prefix in ("", "+", "++"): + for value in ("true", "false"): + with self.subTest(prefix=prefix, value=value): + override = prefix + "resume_from_cache=" + value + args = self.eval_arguments(override, env={"ROLLOUTS_FPATH": "results/existing/resumable.jsonl"}) + self.assertEqual(self.settings(args, "resume_from_cache"), [override]) + self.assertIn("++output_jsonl_fpath=results/existing/resumable.jsonl", args) + + def test_resume_without_fixed_path_uses_restart_timestamp(self): + """Show that enabling resume alone still selects a new output path when the same job restarts.""" + for timestamp in ("20260909_120000", "20260909_180000"): + with self.subTest(timestamp=timestamp): + args = self.eval_arguments("++resume_from_cache=true", env={"TEST_DATE": timestamp}) + self.assertEqual(self.settings(args, "resume_from_cache"), ["++resume_from_cache=true"]) + self.assertIn( + f"++output_jsonl_fpath=results/launcher-test/slurm_job_id_12345/date_{timestamp}.jsonl", args + ) + + def test_explicit_resume_keeps_output_path_across_restarts(self): + """Keep the supplied resume path across requeues and new job IDs, with timestamped logs and W&B names.""" + for job_id in ("12345", "12346"): + for timestamp in ("20260909_120000", "20260909_180000"): + with self.subTest(job_id=job_id, timestamp=timestamp): + args = self.eval_arguments( + "++resume_from_cache=true", + env={ + "ROLLOUTS_FPATH": "results/existing/resumable.jsonl", + "SLURM_JOB_ID": job_id, + "TEST_DATE": timestamp, + }, + ) + self.assertEqual(self.settings(args, "resume_from_cache"), ["++resume_from_cache=true"]) + self.assertIn("++output_jsonl_fpath=results/existing/resumable.jsonl", args) + experiment_name = f"launcher-test/slurm_job_id_{job_id}/date_{timestamp}" + self.assertIn(f"+wandb_name={experiment_name}", args) + self.assertIn(f"+nemo_gym_log_dir=results/{experiment_name}/logs", args) + + def test_prefill_exit_is_detected_during_both_health_checks(self): + """Detect prefill exits during either startup health check, including after request timeouts.""" + _, command = self.generate_commands(env={"VLLM_PD_DEPLOYMENT_MODE": "coupled"}) + stubs = r""" +VLLM_COMMON_ARGS=() +VLLM_PREFILL_ARGS=() +vllm() { command sleep 0.1; return "$TEST_PREFILL_STATUS"; } +curl() { + if [[ "$TEST_WAIT_ROLE" == decode && "$*" == *':8001/health'* ]]; then + return 0 + fi + return "$TEST_HEALTH_STATUS" +} +sleep() { command sleep 0.01; } +hostname() { printf 'node0\n'; } +vllm-router() { printf 'unexpected-router-start\n'; } +""" + for role in ("prefill", "decode"): + for exit_status in (0, 7): + # curl returns 28 for a timed-out request. That must not hide the + # prefill exit status or abort the polling loop prematurely. + for health_status in (1, 28): + with self.subTest(role=role, exit_status=exit_status, health_status=health_status): + status, stdout, stderr = self.run_shell( + stubs + command, + env={ + "TEST_WAIT_ROLE": role, + "TEST_PREFILL_STATUS": str(exit_status), + "TEST_HEALTH_STATUS": str(health_status), + }, + ) + expected_status = exit_status or 1 + self.assertEqual(status, expected_status, stderr) + self.assertNotIn("unexpected-router-start", stdout) + self.assertIn( + f"prefill vLLM process exited while waiting for {role} health (status={expected_status})", + stderr, + ) + + def test_coupled_health_probes_have_timeouts_and_retry_until_ready(self): + """Bound both tiers' health probes and retry timed-out requests until the servers are ready.""" + _, command = self.generate_commands(env={"VLLM_PD_DEPLOYMENT_MODE": "coupled"}) + stubs = r""" +VLLM_COMMON_ARGS=() +VLLM_PREFILL_ARGS=() +vllm() { while true; do command sleep 0.01; done; } +probe_count=0 +curl() { + printf '%s\0' "$@" >&2 + printf '\0' >&2 + probe_count=$((probe_count + 1)) + # Both tiers time out once, then become healthy while prefill stays alive. + if (( probe_count % 2 )); then return 28; fi + return 0 +} +sleep() { :; } +hostname() { printf 'node0\n'; } +vllm-router() { printf 'router-started\n'; kill -TERM "$$"; } +""" + status, stdout, stderr = self.run_shell(stubs + command) + self.assertEqual(status, 143, stderr) + self.assertEqual(stdout, "router-started\n") + probes = [probe.split("\0") for probe in stderr.rstrip("\0").split("\0\0")] + self.assertEqual(len(probes), 4) + self.assertEqual( + [args[-1] for args in probes], + ["http://node0:8001/health"] * 2 + ["http://node4:8002/health"] * 2, + ) + for args in probes: + self.assertIn("--connect-timeout", args) + self.assertEqual(args[args.index("--connect-timeout") + 1], "5") + self.assertIn("--max-time", args) + self.assertEqual(args[args.index("--max-time") + 1], "10") + + def test_healthy_coupled_servers_start_router(self): + """Start the router with the correct prefill and decode URLs once both servers are healthy.""" + _, command = self.generate_commands(env={"VLLM_PD_DEPLOYMENT_MODE": "coupled"}) + stubs = r""" +VLLM_COMMON_ARGS=() +VLLM_PREFILL_ARGS=() +vllm() { while true; do command sleep 0.01; done; } +curl() { return 0; } +hostname() { printf 'node0\n'; } +vllm-router() { printf '%s\0' "$@"; kill -TERM "$$"; } +""" + status, stdout, stderr = self.run_shell(stubs + command) + self.assertEqual(status, 143, stderr) + args = stdout.rstrip("\0").split("\0") + self.assertEqual(args[args.index("--prefill") + 1], "http://node0:8001") + self.assertEqual(args[args.index("--decode") + 1], "http://node4:8002") + + def run_coupled_lifecycle( + self, *, exit_role="", exit_status=0, shutdown_signal="", startup_shutdown=False, exit_before_monitoring=False + ): + _, command = self.generate_commands(env={"VLLM_PD_DEPLOYMENT_MODE": "coupled"}) + stubs = r""" +VLLM_COMMON_ARGS=() +VLLM_PREFILL_ARGS=() +run_service() { + local role=$1 + trap 'printf "%s-stopped\n" "$role"; exit 0' TERM + trap 'exit "$TEST_EXIT_STATUS"' USR1 + printf '%s-started\n' "$role" + touch "$TEST_STATE_DIR/$role-ready" + if [[ "$role" == "$TEST_EXIT_ROLE" ]]; then + # Synchronize on real process startup, not an assumed sleep duration. + while [[ ! -f "$TEST_STATE_DIR/prefill-ready" || ! -f "$TEST_STATE_DIR/router-ready" ]]; do + command sleep 0.01 + done + printf '%s-exited\n' "$role" + return "$TEST_EXIT_STATUS" + fi + if [[ "$role" == router && -n "$TEST_SHUTDOWN_SIGNAL" ]]; then + # $$ remains the supervising shell's PID inside a background function. + kill -s "$TEST_SHUTDOWN_SIGNAL" "$$" + fi + while true; do command sleep 0.01; done +} +vllm() { run_service prefill; } +vllm-router() { run_service router; } +curl() { + [[ -f "$TEST_STATE_DIR/prefill-ready" ]] || return 1 + if [[ "$TEST_STARTUP_SHUTDOWN" == 1 ]]; then + kill -s "$TEST_SHUTDOWN_SIGNAL" "$$" + fi + if [[ "$TEST_EXIT_BEFORE_MONITORING" == 1 && "$*" == *':8002/health'* ]]; then + # Make prefill exit before the final health request returns success. + kill -USR1 "$prefill_pid" + wait "$prefill_pid" || true + fi + return 0 +} +sleep() { command sleep 0.01; } +hostname() { printf 'node0\n'; } +""" + with TemporaryDirectory(prefix="gym-coupled-lifecycle-") as state_dir: + return self.run_shell( + stubs + command, + env={ + "TEST_STATE_DIR": state_dir, + "TEST_EXIT_ROLE": exit_role, + "TEST_EXIT_STATUS": str(exit_status), + "TEST_SHUTDOWN_SIGNAL": shutdown_signal, + "TEST_STARTUP_SHUTDOWN": str(int(startup_shutdown)), + "TEST_EXIT_BEFORE_MONITORING": str(int(exit_before_monitoring)), + }, + ) + + def test_prefill_exit_after_router_start_stops_router(self): + """Treat a prefill exit after router startup as a failure and stop the router.""" + for exit_status in (0, 7): + with self.subTest(exit_status=exit_status): + status, stdout, stderr = self.run_coupled_lifecycle(exit_role="prefill", exit_status=exit_status) + expected_status = exit_status or 1 + self.assertEqual(status, expected_status, stderr) + self.assertIn("router-started\n", stdout) + self.assertIn("prefill-exited\n", stdout) + self.assertIn("router-stopped\n", stdout) + self.assertIn(f"prefill process exited after startup (status={expected_status})", stderr) + + def test_router_exit_after_startup_stops_prefill(self): + """Propagate unexpected router exits as failures and stop prefill.""" + for exit_status in (0, 9): + with self.subTest(exit_status=exit_status): + status, stdout, stderr = self.run_coupled_lifecycle(exit_role="router", exit_status=exit_status) + expected_status = exit_status or 1 + self.assertEqual(status, expected_status, stderr) + self.assertIn("router-exited\n", stdout) + self.assertIn("prefill-stopped\n", stdout) + self.assertIn(f"router process exited after startup (status={expected_status})", stderr) + + def test_coupled_shutdown_signals_stop_both_processes(self): + """Stop both services on SIGTERM or SIGINT while preserving the cancellation exit status.""" + for signal_name, expected_status in (("TERM", 143), ("INT", 130)): + with self.subTest(signal=signal_name): + status, stdout, stderr = self.run_coupled_lifecycle(shutdown_signal=signal_name) + self.assertEqual(status, expected_status, stderr) + self.assertIn("prefill-stopped\n", stdout) + self.assertIn("router-stopped\n", stdout) + self.assertNotIn("ERROR:", stderr) + + def test_coupled_shutdown_during_readiness_stops_prefill(self): + """Stop prefill on cancellation during readiness without starting the router.""" + for signal_name, expected_status in (("TERM", 143), ("INT", 130)): + with self.subTest(signal=signal_name): + status, stdout, stderr = self.run_coupled_lifecycle(shutdown_signal=signal_name, startup_shutdown=True) + self.assertEqual(status, expected_status, stderr) + self.assertIn("prefill-stopped\n", stdout) + self.assertNotIn("router-started\n", stdout) + self.assertNotIn("ERROR:", stderr) + + def test_prefill_exit_before_runtime_monitoring_is_detected(self): + """Catch a prefill exit between the final health response and runtime monitoring.""" + status, _, stderr = self.run_coupled_lifecycle(exit_status=7, exit_before_monitoring=True) + self.assertEqual(status, 7, stderr) + self.assertIn("prefill process exited after startup (status=7)", stderr)