diff --git a/src/cloudai/workloads/ai_dynamo/ai_dynamo.py b/src/cloudai/workloads/ai_dynamo/ai_dynamo.py index 9bd558fe7..30b61bed2 100644 --- a/src/cloudai/workloads/ai_dynamo/ai_dynamo.py +++ b/src/cloudai/workloads/ai_dynamo/ai_dynamo.py @@ -161,6 +161,7 @@ class AIDynamoArgs(BaseModel): model: str = "Qwen/Qwen3-0.6B" backend: Literal["vllm", "sglang", "sglang_dsr1"] = "vllm" + mode: Literal["aggregate", "disaggregate"] = "disaggregate" endpoint: str = Field(default="v1/chat/completions") connector: Optional[str | list[str]] = None @@ -235,12 +236,14 @@ def populate_prefill_decode_args(self) -> "AIDynamoArgs": """Populate prefill/decode args.""" if self.backend.lower() == "vllm": self.prefill_worker.args.model = self.model - self.decode_worker.args.model = self.model + if self.mode == "disaggregate": + self.decode_worker.args.model = self.model elif self.backend.lower() in ["sglang", "sglang_dsr1"]: self.prefill_worker.args.model_path = self.model - self.decode_worker.args.model_path = self.model self.prefill_worker.args.served_model_name = self.model - self.decode_worker.args.served_model_name = self.model + if self.mode == "disaggregate": + self.decode_worker.args.model_path = self.model + self.decode_worker.args.served_model_name = self.model else: raise ValueError(f"Invalid backend: {self.backend}") @@ -577,6 +580,10 @@ def was_run_successful(self, tr: TestRun) -> JobStatusResult: return JobStatusResult(workloads_successful and accuracy_successful) def constraint_check(self, tr: TestRun, system: Optional[System]) -> bool: + if tr.test.cmd_args.dynamo.mode == "aggregate": + logging.info("constraint_check skipped: aggregate mode has no prefill/decode split") + return True + prefill_worker = tr.test.cmd_args.dynamo.prefill_worker decode_worker = tr.test.cmd_args.dynamo.decode_worker diff --git a/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh b/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh index 18ce737ee..ed48b2d64 100644 --- a/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh +++ b/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh @@ -307,7 +307,11 @@ _apply_connector_settings() { _patch_dynamo_args() { if [[ -z "${dynamo_args["frontend-node"]}" ]]; then - dynamo_args["frontend-node"]=$(echo "${decode_config["node-list"]}" | cut -d',' -f1) + if [[ -n "${decode_config["node-list"]:-}" ]]; then + dynamo_args["frontend-node"]=$(echo "${decode_config["node-list"]}" | cut -d',' -f1) + else + dynamo_args["frontend-node"]=$(echo "${prefill_config["node-list"]}" | cut -d',' -f1) + fi fi dynamo_args["url"]="http://${dynamo_args["frontend-node"]}:${dynamo_args["port"]}" @@ -342,13 +346,20 @@ _compute_worker_allocation_vllm() { fi prefill_config["gpus-per-worker"]=$(( prefill_args["--tensor-parallel-size"] * prefill_args["--pipeline-parallel-size"] )) - decode_config["gpus-per-worker"]=$(( decode_args["--tensor-parallel-size"] * decode_args["--pipeline-parallel-size"] )) - if [[ ${prefill_config["gpus-per-worker"]} -eq 0 ]] || [[ ${decode_config["gpus-per-worker"]} -eq 0 ]]; then + if [[ ${prefill_config["gpus-per-worker"]} -eq 0 ]]; then log "ERROR: Invalid TP/PP configuration" exit 1 fi + if [[ "${decode_config["num-nodes"]:-0}" -gt 0 ]]; then + decode_config["gpus-per-worker"]=$(( decode_args["--tensor-parallel-size"] * decode_args["--pipeline-parallel-size"] )) + if [[ ${decode_config["gpus-per-worker"]} -eq 0 ]]; then + log "ERROR: Invalid decode TP/PP configuration" + exit 1 + fi + fi + decode_config["gpu-offset"]=0 prefill_config["gpu-offset"]=0 @@ -362,16 +373,21 @@ _compute_worker_allocation_vllm() { prefill_config["workers-per-node"]=1 prefill_config["gpu-offset"]=${decode_config["gpus-per-worker"]} else - if [[ "${prefill_config["multiple-workers-per-node"]}" != "true" ]]; then + if [[ "${prefill_config["multiple-workers-per-node"],,}" != "true" ]]; then prefill_config["gpus-per-worker"]=$num_gpus fi - if [[ "${decode_config["multiple-workers-per-node"]}" != "true" ]]; then - decode_config["gpus-per-worker"]=$num_gpus - fi - prefill_config["workers-per-node"]=$(( num_gpus / prefill_config["gpus-per-worker"] )) - decode_config["workers-per-node"]=$(( num_gpus / decode_config["gpus-per-worker"] )) + + if [[ "${decode_config["num-nodes"]:-0}" -gt 0 ]]; then + if [[ "${decode_config["multiple-workers-per-node"],,}" != "true" ]]; then + decode_config["gpus-per-worker"]=$num_gpus + fi + decode_config["workers-per-node"]=$(( num_gpus / decode_config["gpus-per-worker"] )) + else + decode_config["gpus-per-worker"]=0 + decode_config["workers-per-node"]=0 + fi fi log "DECODE: num GPUs: $num_gpus, GPUs per worker: ${decode_config["gpus-per-worker"]}" @@ -1124,7 +1140,7 @@ _resolve_aiperf_server_metrics_urls() { done done - if [[ "${dynamo_args["dcgm-exporter-enabled"]}" == "True" || "${dynamo_args["dcgm-exporter-enabled"]}" == "true" ]]; then + if [[ "${dynamo_args["dcgm-exporter-enabled"],,}" == "true" ]]; then local dcgm_nodes="${decode_config["node-list"]:-},${prefill_config["node-list"]:-}" for node in $dcgm_nodes; do [[ -z "$node" ]] && continue @@ -1393,12 +1409,14 @@ function main() launch_etcd & launch_nats & wait_for_etcd - launch_ingress - if _is_sglang_dsr1; then - launch_sgl_http_server - fi fi + # Workers launch BEFORE the ingress: launch_ingress blocks in + # wait_for_router, and the router only becomes ready once a worker + # registers — on a combined frontend+worker node the old order serialized + # the whole ROUTER_START_TIMEOUT (120 s of failing readiness curls) in + # front of every worker start. Workers only need etcd/nats (waited above) + # and the lmcache config from setup_lmcache; they never talk to the router. if _is_decode_node; then log "Node ID: $SLURM_NODEID, Role: decode" log_node_role "$(_current_node_name)" "decode" @@ -1412,6 +1430,10 @@ function main() fi if _is_frontend_node; then + launch_ingress + if _is_sglang_dsr1; then + launch_sgl_http_server + fi sleep 10 launch_workloads & diff --git a/src/cloudai/workloads/ai_dynamo/slurm_command_gen_strategy.py b/src/cloudai/workloads/ai_dynamo/slurm_command_gen_strategy.py index c7370b949..6b200f194 100644 --- a/src/cloudai/workloads/ai_dynamo/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/ai_dynamo/slurm_command_gen_strategy.py @@ -400,13 +400,17 @@ def _gen_script_args(self, td: AIDynamoTestDefinition) -> List[str]: args.append('--dynamo-dcgm-exporter-enabled "True"') args.append(f'--dynamo-dcgm-exporter-port "{td.cmd_args.dynamo.dcgm_exporter.port}"') + is_aggregate = td.cmd_args.dynamo.mode == "aggregate" + if td.cmd_args.dynamo.prefill_worker: args.extend(self._get_nested_toml_args(td.cmd_args.dynamo.prefill_worker, "--prefill-", exclude=["nodes"])) if td.cmd_args.dynamo.prefill_worker.nodes: args.append(f"--prefill-node-list {shlex.quote(td.cmd_args.dynamo.prefill_worker.nodes)}") - args.extend(self._get_nested_toml_args(td.cmd_args.dynamo.decode_worker, "--decode-", exclude=["nodes"])) - if td.cmd_args.dynamo.decode_worker.nodes: - args.append(f"--decode-node-list {shlex.quote(td.cmd_args.dynamo.decode_worker.nodes)}") + + if not is_aggregate: + args.extend(self._get_nested_toml_args(td.cmd_args.dynamo.decode_worker, "--decode-", exclude=["nodes"])) + if td.cmd_args.dynamo.decode_worker.nodes: + args.append(f"--decode-node-list {shlex.quote(td.cmd_args.dynamo.decode_worker.nodes)}") args.extend(self._get_nested_toml_args(td.cmd_args.genai_perf, "--genai_perf-")) if aiperf_script: @@ -608,12 +612,16 @@ def get_cached_nodes_spec(self) -> tuple[int, list[str]]: if cache_key in self._node_spec_cache: return self._node_spec_cache[cache_key] + is_aggregate = self.td.cmd_args.dynamo.mode == "aggregate" + prefill_n, prefill_nodes = 0, "" if self.td.cmd_args.dynamo.prefill_worker: prefill_n = cast(int, self.td.cmd_args.dynamo.prefill_worker.num_nodes) prefill_nodes = self.td.cmd_args.dynamo.prefill_worker.nodes - decode_n = self.td.cmd_args.dynamo.decode_worker.num_nodes - decode_nodes = self.td.cmd_args.dynamo.decode_worker.nodes + + # In aggregate mode there is no separate decode worker; all nodes run the prefill worker. + decode_n = 0 if is_aggregate else cast(int, self.td.cmd_args.dynamo.decode_worker.num_nodes) + decode_nodes = None if is_aggregate else self.td.cmd_args.dynamo.decode_worker.nodes assert isinstance(prefill_n, int), "prefill_worker.num_nodes must be an integer" assert isinstance(decode_n, int), "decode_worker.num_nodes must be an integer" @@ -649,7 +657,8 @@ def get_cached_nodes_spec(self) -> tuple[int, list[str]]: if prefill_nodes or decode_nodes: self._validate_worker_nodes(node_list, prefill_nodes, prefill_n, "prefill") - self._validate_worker_nodes(node_list, decode_nodes, decode_n, "decode") + if not is_aggregate: + self._validate_worker_nodes(node_list, decode_nodes, decode_n, "decode") if self._worker_nodes_overlap(prefill_nodes, decode_nodes): unique_worker_nodes = self._unique_nodes( self._split_node_list(prefill_nodes) + self._split_node_list(decode_nodes)