Merge pull request #745 from NVIDIA/am/dynamo-slurm

amaslenn · web-flow · commit 8e26c01948e4 · 2025-12-10T17:08:04.000+01:00
Simplify Dynamo slurm configuration
diff --git a/conf/experimental/ai_dynamo/test/vllm.toml b/conf/experimental/ai_dynamo/test/vllm.toml
@@ -28,9 +28,6 @@ docker_image_url = "nvcr.io/nvidia/ai-dynamo/vllm-runtime:0.6.1.post1"
   prefill-cmd = 'python3 -m dynamo.vllm --is-prefill-worker'
   decode-cmd = 'python3 -m dynamo.vllm'
 
-    [cmd_args.dynamo.prefill_worker]
-    pipeline-parallel-size = 1
-
     [cmd_args.dynamo.decode_worker]
     pipeline-parallel-size = 1
 
@@ -42,13 +39,14 @@ docker_image_url = "nvcr.io/nvidia/ai-dynamo/vllm-runtime:0.6.1.post1"
   output-tokens-mean = 500
   output-tokens-stddev = 0
   random-seed = 123
-  request-count = 128
+  request-count = 50
   synthetic-input-tokens-mean = 300
   synthetic-input-tokens-stddev = 0
-  warmup-request-count = 10
-  concurrency = 1
+  warmup-request-count = 5
+  concurrency = 2
   extra-args = "--streaming -- -v --async"
 
 [extra_env_vars]
 UCX_LOG_LEVEL = "warn"
 UCX_TLS = "cuda_copy,rc_x"
+DYNAMO_NODELIST = "$(scontrol show hostname $SLURM_JOB_NODELIST | tr -s '\\n' ',')"
diff --git a/conf/experimental/ai_dynamo/test_scenario/vllm_slurm.toml b/conf/experimental/ai_dynamo/test_scenario/vllm_slurm.toml
@@ -15,40 +15,36 @@
 # limitations under the License.
 
 name = "dynamo-vllm-slurm"
+job_status_check = false
 
 [[Tests]]
-id = "qwen3-0.6B"
+id = "test.disagg.single-node"
 test_name = "vLLM-Qwen3-0.6B"
-num_nodes = 4
-time_limit = "00:20:00"
+num_nodes = 2                  # 1 prefill node + 1 decode node
+time_limit = "00:10:00"
 
-  [Tests.cmd_args]
+  [Tests.cmd_args.dynamo.prefill_worker]
+  num-nodes = 1
+  tensor-parallel-size = 4
+  pipeline-parallel-size = 1
 
-    [Tests.cmd_args.dynamo]
-    decode-initialized-regex = 'VllmWorker.*has.been.initialized'
-    etcd-cmd = "etcd --log-level debug"
-    etcd-port = 2379
-    genai-perf-cmd = 'genai-perf profile'
-    ingress-cmd = "python -m dynamo.frontend --router-mode kv"
-    nats-cmd = "nats-server -js"
-    nats-port = 4222
-    node-setup-cmd = "apt-get update -o APT::Sandbox::User=root && apt-get install -y curl libibverbs1 rdma-core ibverbs-utils libibumad3 libnuma1 librdmacm1 ibverbs-providers; /usr/local/ucx/bin/ucx_info -d |grep Transport | sort -u;"
-    port = 8787
-    prefill-initialized-regex = 'VllmWorker.*has.been.initialized'
+  [Tests.cmd_args.dynamo.decode_worker]
+  num-nodes = 1
+  tensor-parallel-size = 4
+  pipeline-parallel-size = 1
 
-      [Tests.cmd_args.dynamo.prefill_worker]
-      gpu-memory-utilization = 0.90
-      max_model_len = 19280
-      num-nodes = 2
-      tensor-parallel-size = 4
+[[Tests]]
+id = "test.disagg.multinode"
+test_name = "vLLM-Qwen3-0.6B"
+num_nodes = 4                 # 2 prefill nodes + 2 decode nodes
+time_limit = "00:10:00"
 
-      [Tests.cmd_args.dynamo.decode_worker]
-      gpu-memory-utilization = 0.90
-      max_model_len = 19280
-      num-nodes = 2
-      tensor-parallel-size = 4
+  [Tests.cmd_args.dynamo.prefill_worker]
+  num-nodes = 2
+  tensor-parallel-size = 4
+  pipeline-parallel-size = 1
 
-  [Tests.extra_env_vars]
-  UCX_LOG_LEVEL = "warn"
-  UCX_TLS = "cuda_copy,rc_x"
-  DYNAMO_NODELIST = "$(scontrol show hostname $SLURM_JOB_NODELIST | tr -s '\\n' ',')"
+  [Tests.cmd_args.dynamo.decode_worker]
+  num-nodes = 2
+  tensor-parallel-size = 4
+  pipeline-parallel-size = 1
diff --git a/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh b/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh
@@ -33,8 +33,6 @@ dynamo_args["ingress-cmd"]="python -m dynamo.frontend --router-mode kv"
 dynamo_args["port"]=8080
 dynamo_args["endpoint"]="v1/chat/completions"
 dynamo_args["model"]="deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
-dynamo_args["etcd-cmd"]="etcd --log-level debug"
-dynamo_args["nats-cmd"]="nats-server -js"
 dynamo_args["etcd-port"]=2379
 dynamo_args["nats-port"]=4222
 dynamo_args["workspace-path"]="/workspace"
@@ -47,8 +45,12 @@ dynamo_args["tp-arg-name"]="tensor-parallel-size"
 dynamo_args["pp-arg-name"]="pipeline-parallel-size"
 dynamo_args["multiple-prefill-workers-per-node"]="true"
 dynamo_args["multiple-decode-workers-per-node"]="true"
-dynamo_args["prefill-initialized-regex"]="prefill.*initialized"
-dynamo_args["decode-initialized-regex"]="decode.*initialized"
+dynamo_args["prefill-initialized-regex"]="Worker.*has.been.initialized"
+dynamo_args["decode-initialized-regex"]="Worker.*has.been.initialized"
+
+dynamo_args["etcd-cmd"]="etcd --log-level debug"
+dynamo_args["nats-cmd"]="nats-server -js"
+dynamo_args["genai-perf-cmd"]="genai-perf profile"
 
 # sglang-specific optional ports. Ignored by vllm.
 dynamo_args["sgl-http-port"]=9001
@@ -310,15 +312,21 @@ _compute_worker_allocation_vllm() {
     dynamo_args["decode-gpus-per-worker"]=$num_gpus
   fi
 
+  log "DECODE: num GPUs: $num_gpus, GPUs per worker: ${dynamo_args["decode-gpus-per-worker"]}"
+  log "PREFILL: num GPUs: $num_gpus, GPUs per worker: ${dynamo_args["prefill-gpus-per-worker"]}"
   dynamo_args["prefill-workers-per-node"]=$(( num_gpus / dynamo_args["prefill-gpus-per-worker"] ))
   dynamo_args["decode-workers-per-node"]=$(( num_gpus / dynamo_args["decode-gpus-per-worker"] ))
+  log "DECODE: workers per node: ${dynamo_args["decode-workers-per-node"]}"
+  log "PREFILL: workers per node: ${dynamo_args["prefill-workers-per-node"]}"
 
   if [[ -n "${prefill_args["--num-nodes"]}" ]]; then
     dynamo_args["num-prefill-nodes"]=${prefill_args["--num-nodes"]}
   fi
   if [[ -n "${decode_args["--num-nodes"]}" ]]; then
     dynamo_args["num-decode-nodes"]=${decode_args["--num-nodes"]}
   fi
+  log "NUM PREFILL NODES: ${dynamo_args["num-prefill-nodes"]}"
+  log "NUM DECODE NODES: ${dynamo_args["num-decode-nodes"]}"
 }
 
 _compute_worker_allocation() {
@@ -597,7 +605,7 @@ validate_environment() {
 
 function launch_etcd()
 {
-  log "Launching etcd"
+  log "Launching etcd with cmd: ${dynamo_args["etcd-cmd"]} --listen-client-urls http://0.0.0.0:${dynamo_args["etcd-port"]} --advertise-client-urls http://0.0.0.0:${dynamo_args["etcd-port"]}"
   ${dynamo_args["etcd-cmd"]} \
     --listen-client-urls http://0.0.0.0:${dynamo_args["etcd-port"]} \
     --advertise-client-urls http://0.0.0.0:${dynamo_args["etcd-port"]} \
@@ -606,7 +614,7 @@ function launch_etcd()
 
 function launch_nats()
 {
-  log "Launching nats"
+  log "Launching nats with cmd: ${dynamo_args["nats-cmd"]} -p ${dynamo_args["nats-port"]}"
   ${dynamo_args["nats-cmd"]} -p ${dynamo_args["nats-port"]} > ${RESULTS_DIR}/nats.log 2>&1
 }
 
@@ -633,12 +641,14 @@ function launch_decode()
   wait_for_etcd
 
   local workers_per_node=${dynamo_args["decode-workers-per-node"]}
+  log "Using workers per node: $workers_per_node"
 
   for i in $(seq 0 $(( $workers_per_node - 1 ))); do
     local gpu_list=$(_gpu_list_for_worker "${dynamo_args["decode-gpus-per-worker"]}" "$i")
     local log_file=$(_log_file_for_worker "decode" "$i")
 
     log "Launching decode worker $i on GPUs $gpu_list"
+    log "Decode cmd: ${dynamo_args["decode-cmd"]} $(array_to_args decode_args) ${decode_args["--extra-args"]}"
     CUDA_VISIBLE_DEVICES=$gpu_list \
       ${dynamo_args["decode-cmd"]} \
       $(array_to_args decode_args) ${decode_args["--extra-args"]} > $log_file 2>&1 &
@@ -665,6 +675,7 @@ function launch_prefill()
     local log_file=$(_log_file_for_worker "prefill" "$i")
 
     log "Launching prefill worker $i on GPUs $gpu_list"
+    log "Prefill cmd: ${dynamo_args["prefill-cmd"]} $(array_to_args prefill_args) ${prefill_args["--extra-args"]}"
     CUDA_VISIBLE_DEVICES=$gpu_list \
       ${dynamo_args["prefill-cmd"]} \
       $(array_to_args prefill_args) ${prefill_args["--extra-args"]} > $log_file 2>&1 &
@@ -680,11 +691,12 @@ function wait_for_dynamo_frontend()
     local have_prefill=$(_count_initialized_prefill)
     local have_decode=$(_count_initialized_decode)
 
+    log "Initialized: prefill ${have_prefill}/${want_prefill}; decode ${have_decode}/${want_decode}"
+
     if [[ $have_prefill -ge $want_prefill && $have_decode -ge $want_decode ]]; then
       break
     fi
 
-    log "Initialized: prefill ${have_prefill}/${want_prefill}; decode ${have_decode}/${want_decode}"
     exit_on_error
     sleep 30
   done
@@ -710,7 +722,7 @@ function launch_genai_perf()
   echo "Response: $resp"
 
   local genai_perf_arguments=$(array_to_args genai_perf_args)
-  log "Launching genai-perf with args: $genai_perf_arguments ${genai_perf_args["--extra-args"]}"
+  log "Launching genai-perf with cmd: ${dynamo_args["genai-perf-cmd"]} $genai_perf_arguments ${genai_perf_args["--extra-args"]}"
 
   ${dynamo_args["genai-perf-cmd"]} ${genai_perf_arguments} ${genai_perf_args["--extra-args"]} > ${RESULTS_DIR}/genai_perf.log 2>&1
 
