diff --git a/tools/dynamo/slurm_smoke/README.md b/tools/dynamo/slurm_smoke/README.md new file mode 100644 index 0000000000..c69d891df4 --- /dev/null +++ b/tools/dynamo/slurm_smoke/README.md @@ -0,0 +1,87 @@ +# Dynamo SLURM Smoke Test + +This example runs a single-node smoke test with Dynamo serving inference and +prime-rl training against that Dynamo endpoint. + +It is intended to mirror the local/Kubernetes Dynamo smoke flow in +`tools/dynamo/` and `k8s/dynamo-deploy/`, but under one SLURM allocation: + +- GPU 0: Dynamo frontend + vLLM worker, using `--discovery-backend file` +- GPU 1: prime-rl orchestrator + trainer + +`--discovery-backend file` keeps the smoke test self-contained on one node and +does not require etcd or NATS. + +## Prerequisites + +- prime-rl is cloned on a filesystem visible to the SLURM compute node. +- The prime-rl virtualenv exists at `/.venv`. +- Dynamo with vLLM support is installed either in the same virtualenv or in a + separate virtualenv pointed to by `DYNAMO_VENV`. +- The SLURM node has at least two GPUs. + +If Dynamo is installed separately: + +```bash +export DYNAMO_VENV=/shared/dynamo/.venv +``` + +## Dry Run + +Render the resolved config and sbatch script without submitting: + +```bash +uv run rl @ tools/dynamo/slurm_smoke/smoke_rl.toml \ + --slurm.project-dir /shared/prime-rl \ + --slurm.partition \ + --dry-run +``` + +The generated script is written to: + +```text +outputs/dynamo-slurm-smoke/rl.sbatch +``` + +## Submit + +```bash +uv run rl @ tools/dynamo/slurm_smoke/smoke_rl.toml \ + --slurm.project-dir /shared/prime-rl \ + --slurm.partition +``` + +Useful overrides: + +```bash +export DYNAMO_MODEL=PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT +export DYNAMO_SERVED_MODEL_NAME=$DYNAMO_MODEL +export DYNAMO_GPU=0 +export PRIME_RL_GPU=1 +export DYN_HTTP_PORT=8000 +export DYNAMO_MAX_MODEL_LEN=2048 +export DYNAMO_MAX_NUM_SEQS=32 +``` + +Additional Dynamo worker flags can be passed with `DYNAMO_EXTRA_ARGS`. + +## Logs + +After submission, logs are under the run output directory: + +```text +outputs/dynamo-slurm-smoke/job_.log +outputs/dynamo-slurm-smoke/logs/dynamo/frontend.log +outputs/dynamo-slurm-smoke/logs/dynamo/vllm.log +outputs/dynamo-slurm-smoke/logs/orchestrator.log +outputs/dynamo-slurm-smoke/logs/trainer.log +``` + +The prime-rl config disables `use_token_client` and points both inference and +admin traffic at the Dynamo frontend: + +```toml +[orchestrator.client] +base_url = ["http://127.0.0.1:8000/v1"] +admin_base_url = ["http://127.0.0.1:8000/v1/rl"] +``` diff --git a/tools/dynamo/slurm_smoke/dynamo_single_node_rl.sbatch.j2 b/tools/dynamo/slurm_smoke/dynamo_single_node_rl.sbatch.j2 new file mode 100644 index 0000000000..c482ba8786 --- /dev/null +++ b/tools/dynamo/slurm_smoke/dynamo_single_node_rl.sbatch.j2 @@ -0,0 +1,139 @@ +#!/bin/bash + +#SBATCH --job-name={{ job_name }} +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:{{ gpus_per_node }} +#SBATCH --exclusive +#SBATCH --partition={{ partition }} +{%- if account %} +#SBATCH --account={{ account }} +{%- endif %} +{%- if time %} +#SBATCH --time={{ time }} +{%- endif %} +{%- if nodelist %} +#SBATCH --nodelist={{ nodelist }} +{%- endif %} +{%- if exclude %} +#SBATCH --exclude={{ exclude }} +{%- endif %} +#SBATCH --output={{ output_dir }}/job_%j.log +#SBATCH --error={{ output_dir }}/job_%j.log + +set -euo pipefail + +export PROJECT_DIR={{ project_dir }} +export OUTPUT_DIR={{ output_dir }} +export CONFIG_PATH={{ config_path }} + +mkdir -p "$OUTPUT_DIR/logs/dynamo" + +# Runtime knobs. Override these in the submit shell or your .env file. +export DYNAMO_VENV="${DYNAMO_VENV:-$PROJECT_DIR/.venv}" +export DYNAMO_MODEL="${DYNAMO_MODEL:-PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT}" +export DYNAMO_SERVED_MODEL_NAME="${DYNAMO_SERVED_MODEL_NAME:-$DYNAMO_MODEL}" +export DYNAMO_GPU="${DYNAMO_GPU:-0}" +export PRIME_RL_GPU="${PRIME_RL_GPU:-1}" +export DYN_HTTP_PORT="${DYN_HTTP_PORT:-8000}" +export DYN_SYSTEM_PORT="${DYN_SYSTEM_PORT:-8081}" +export DYN_ENABLE_RL="${DYN_ENABLE_RL:-true}" +export DYNAMO_MAX_MODEL_LEN="${DYNAMO_MAX_MODEL_LEN:-2048}" +export DYNAMO_MAX_NUM_SEQS="${DYNAMO_MAX_NUM_SEQS:-32}" +export DYNAMO_STARTUP_ATTEMPTS="${DYNAMO_STARTUP_ATTEMPTS:-180}" + +cleanup() { + set +e + jobs -pr | xargs -r kill +} +trap cleanup EXIT + +wait_for_url() { + local url="$1" + local name="$2" + local attempts="$DYNAMO_STARTUP_ATTEMPTS" + + for ((i = 1; i <= attempts; i++)); do + if curl -fsS "$url" >/dev/null 2>&1; then + echo "[dynamo-slurm] $name is ready: $url" + return 0 + fi + + if [[ -n "${FRONTEND_PID:-}" ]] && ! kill -0 "$FRONTEND_PID" 2>/dev/null; then + echo "[dynamo-slurm] frontend exited while waiting for $name" + tail -n 80 "$OUTPUT_DIR/logs/dynamo/frontend.log" || true + return 1 + fi + if [[ -n "${WORKER_PID:-}" ]] && ! kill -0 "$WORKER_PID" 2>/dev/null; then + echo "[dynamo-slurm] worker exited while waiting for $name" + tail -n 80 "$OUTPUT_DIR/logs/dynamo/vllm.log" || true + return 1 + fi + + sleep 5 + done + + echo "[dynamo-slurm] timed out waiting for $name: $url" + tail -n 80 "$OUTPUT_DIR/logs/dynamo/frontend.log" || true + tail -n 80 "$OUTPUT_DIR/logs/dynamo/vllm.log" || true + return 1 +} + +cd "$PROJECT_DIR" +[ -f .env ] && source .env +source .venv/bin/activate +uv sync --all-extras + +{% if pre_run_command %} +# Pre-run command +{{ pre_run_command }} +{% endif %} + +echo "[dynamo-slurm] cleaning stale local processes and IPC files" +pkill -9 -f "python.*prime_rl" 2>/dev/null || true +pkill -9 -f "torchrun" 2>/dev/null || true +pkill -9 -f "dynamo.frontend" 2>/dev/null || true +pkill -9 -f "dynamo.vllm" 2>/dev/null || true +pkill -9 -f "vllm" 2>/dev/null || true +rm -rf /dev/shm/vllm-* /dev/shm/vllm_* /tmp/vllm-* /tmp/vllm_* /tmp/torch-* /tmp/torchelastic_* 2>/dev/null || true + +echo "[dynamo-slurm] launching Dynamo frontend on port $DYN_HTTP_PORT" +( + source "$DYNAMO_VENV/bin/activate" + export DYN_ENABLE_RL + export DYN_HTTP_PORT + python -m dynamo.frontend \ + --discovery-backend file \ + > "$OUTPUT_DIR/logs/dynamo/frontend.log" 2>&1 +) & +FRONTEND_PID=$! + +sleep 5 + +echo "[dynamo-slurm] launching Dynamo vLLM worker on GPU $DYNAMO_GPU model=$DYNAMO_MODEL" +( + source "$DYNAMO_VENV/bin/activate" + export CUDA_VISIBLE_DEVICES="$DYNAMO_GPU" + export DYN_ENABLE_RL + export DYN_SYSTEM_PORT + export VLLM_WORKER_MULTIPROC_METHOD=spawn + python -m dynamo.vllm \ + --model "$DYNAMO_MODEL" \ + --served-model-name "$DYNAMO_SERVED_MODEL_NAME" \ + --discovery-backend file \ + --kv-events-config '{"enable_kv_cache_events": false}' \ + --enforce-eager \ + --max-model-len "$DYNAMO_MAX_MODEL_LEN" \ + --max-num-seqs "$DYNAMO_MAX_NUM_SEQS" \ + ${DYNAMO_EXTRA_ARGS:-} \ + > "$OUTPUT_DIR/logs/dynamo/vllm.log" 2>&1 +) & +WORKER_PID=$! + +wait_for_url "http://127.0.0.1:${DYN_HTTP_PORT}/health" "Dynamo frontend" +wait_for_url "http://127.0.0.1:${DYN_HTTP_PORT}/v1/models" "Dynamo model registry" + +echo "[dynamo-slurm] starting prime-rl on GPU $PRIME_RL_GPU" +export BENCH_API_KEY="${BENCH_API_KEY:-EMPTY}" +export CUDA_VISIBLE_DEVICES="$PRIME_RL_GPU" +uv run rl @ "$CONFIG_PATH" diff --git a/tools/dynamo/slurm_smoke/smoke_rl.toml b/tools/dynamo/slurm_smoke/smoke_rl.toml new file mode 100644 index 0000000000..e241e5fe11 --- /dev/null +++ b/tools/dynamo/slurm_smoke/smoke_rl.toml @@ -0,0 +1,54 @@ +# Prime-RL + Dynamo single-node SLURM smoke test. +# +# Usage: +# uv run rl @ tools/dynamo/slurm_smoke/smoke_rl.toml --dry-run +# uv run rl @ tools/dynamo/slurm_smoke/smoke_rl.toml +# +# The custom SLURM template launches Dynamo on GPU 0, waits for it to become +# ready, then runs prime-rl against that Dynamo endpoint on GPU 1. + +output_dir = "outputs/dynamo-slurm-smoke" +max_steps = 5 +seq_len = 512 + +[model] +name = "PrimeIntellect/Qwen3-0.6B-Reverse-Text-SFT" + +[wandb] +offline = true +shared = false +project = "prime-rl-dynamo-slurm" +name = "dynamo-slurm-smoke" + +[deployment] +type = "single_node" +gpus_per_node = 2 +num_train_gpus = 1 +num_infer_gpus = 0 + +[slurm] +job_name = "dynamo-slurm-smoke" +time = "01:00:00" +template_path = "tools/dynamo/slurm_smoke/dynamo_single_node_rl.sbatch.j2" + +[orchestrator] +batch_size = 16 +rollouts_per_example = 4 +use_token_client = false + +[orchestrator.train.sampling] +max_completion_tokens = 64 + +[[orchestrator.train.env]] +id = "reverse-text" + +[orchestrator.client] +base_url = ["http://127.0.0.1:8000/v1"] +admin_base_url = ["http://127.0.0.1:8000/v1/rl"] +api_key_var = "BENCH_API_KEY" +skip_model_check = true + +[trainer.optim] +lr = 3e-6 + +[ckpt]