diff --git a/src/cloudai/workloads/ai_dynamo/ai_dynamo.py b/src/cloudai/workloads/ai_dynamo/ai_dynamo.py index d3726fda1..ff95c6b3d 100644 --- a/src/cloudai/workloads/ai_dynamo/ai_dynamo.py +++ b/src/cloudai/workloads/ai_dynamo/ai_dynamo.py @@ -162,9 +162,17 @@ class AIDynamoArgs(BaseModel): model: str = "Qwen/Qwen3-0.6B" backend: Literal["vllm", "sglang", "sglang_dsr1"] = "vllm" + mode: Literal["aggregate", "disaggregated"] = "disaggregated" endpoint: str = Field(default="v1/chat/completions") connector: Optional[str | list[str]] = None + @field_validator("mode", mode="before") + @classmethod + def normalize_mode(cls, v: str) -> str: + if v == "disaggregate": + return "disaggregated" + return v + @field_validator("connector", mode="before") @classmethod def validate_connector(cls, v: str | list[str] | None) -> str | list[str] | None: @@ -234,14 +242,21 @@ def validate_connector(cls, v: str | list[str] | None) -> str | list[str] | None @model_validator(mode="after") def populate_prefill_decode_args(self) -> "AIDynamoArgs": """Populate prefill/decode args.""" + if self.mode == "aggregate" and self.backend.lower() != "vllm": + raise ValueError("AI Dynamo aggregate mode is currently supported only for the vLLM backend") + if self.backend.lower() == "vllm": self.prefill_worker.args.model = self.model - self.decode_worker.args.model = self.model + if self.mode == "aggregate": + self.prefill_worker.cmd = "python3 -m dynamo.vllm" + else: + 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 == "disaggregated": + 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}") @@ -593,19 +608,20 @@ 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: + is_aggregate = tr.test.cmd_args.dynamo.mode == "aggregate" prefill_worker = tr.test.cmd_args.dynamo.prefill_worker decode_worker = tr.test.cmd_args.dynamo.decode_worker prefill_tp = prefill_worker.args.tensor_parallel_size prefill_pp = prefill_worker.args.pipeline_parallel_size - decode_tp = decode_worker.args.tensor_parallel_size decode_pp = decode_worker.args.pipeline_parallel_size - if self.constraints.prefill_tp_le_decode_tp and prefill_tp > decode_tp: - logging.info("constraint_check failed for: prefill_tp_le_decode_tp") - return False - logging.info("constraint_check passed for: prefill_tp_le_decode_tp") + if not is_aggregate: + if self.constraints.prefill_tp_le_decode_tp and prefill_tp > decode_tp: + logging.info("constraint_check failed for: prefill_tp_le_decode_tp") + return False + logging.info("constraint_check passed for: prefill_tp_le_decode_tp") gpus_per_node = 0 slurm_system = cast(SlurmSystem, system) @@ -615,12 +631,24 @@ def constraint_check(self, tr: TestRun, system: Optional[System]) -> bool: if ( gpus_per_node > 0 and self.constraints.tp_times_pp_le_gpus_per_node - and (prefill_tp * prefill_pp > gpus_per_node or decode_tp * decode_pp > gpus_per_node) + and prefill_tp * prefill_pp > gpus_per_node + ): + logging.info("constraint_check failed for: tp_times_pp_le_gpus_per_node") + return False + if ( + not is_aggregate + and gpus_per_node > 0 + and self.constraints.tp_times_pp_le_gpus_per_node + and decode_tp * decode_pp > gpus_per_node ): logging.info("constraint_check failed for: tp_times_pp_le_gpus_per_node") return False logging.info("constraint_check passed for: tp_times_pp_le_gpus_per_node") + if is_aggregate: + logging.info("constraint_check skipped split-worker checks: aggregate mode has no decode worker") + return True + role_total_nodes = int(prefill_worker.num_nodes) + int(decode_worker.num_nodes) prefill_nodes = set(prefill_worker.nodes.split(",")) if prefill_worker.nodes else set() decode_nodes = set(decode_worker.nodes.split(",")) if decode_worker.nodes else set() diff --git a/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh b/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh index f0612b2db..27a722b5f 100644 --- a/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh +++ b/src/cloudai/workloads/ai_dynamo/ai_dynamo.sh @@ -43,6 +43,7 @@ SHARED_NODE_DISAGG="false" declare -A dynamo_args dynamo_args["backend"]="vllm" +dynamo_args["mode"]="disaggregated" dynamo_args["node-setup-cmd"]="" dynamo_args["ingress-cmd"]="python -m dynamo.frontend --router-mode kv" dynamo_args["port"]=$((8080 + SLURM_JOBID % 100)) @@ -72,6 +73,7 @@ function log() _is_vllm() { [[ "${dynamo_args["backend"]}" == "vllm" ]]; } _is_sglang() { [[ "${dynamo_args["backend"]}" == "sglang" ]]; } _is_sglang_dsr1() { [[ "${dynamo_args["backend"]}" == "sglang_dsr1" ]]; } +_is_aggregate() { [[ "${dynamo_args["mode"]}" == "aggregate" ]]; } _csv_len() { grep -oE '[^,]+' <<< "$1" | wc -l; } @@ -211,6 +213,45 @@ _parse_cli_pairs() { done } +_normalize_runtime_mode() { + if [[ "${dynamo_args["mode"]}" == "disaggregate" ]]; then + dynamo_args["mode"]="disaggregated" + fi + + if [[ "${dynamo_args["mode"]}" != "aggregate" && "${dynamo_args["mode"]}" != "disaggregated" ]]; then + log "ERROR: unsupported dynamo mode '${dynamo_args["mode"]}'" + exit 1 + fi + + if _is_aggregate && ! _is_vllm; then + log "ERROR: aggregate mode is currently supported only for the vLLM backend" + exit 1 + fi +} + +_init_role_defaults() { + prefill_config["num-nodes"]="${prefill_config["num-nodes"]:-1}" + prefill_config["node-list"]="${prefill_config["node-list"]:-}" + prefill_config["multiple-workers-per-node"]="${prefill_config["multiple-workers-per-node"]:-false}" + prefill_config["extra-args"]="${prefill_config["extra-args"]:-}" + prefill_args["--tensor-parallel-size"]="${prefill_args["--tensor-parallel-size"]:-1}" + prefill_args["--pipeline-parallel-size"]="${prefill_args["--pipeline-parallel-size"]:-1}" + + if _is_aggregate; then + decode_config["num-nodes"]=0 + decode_config["node-list"]="" + decode_config["multiple-workers-per-node"]="false" + decode_config["extra-args"]="" + else + decode_config["num-nodes"]="${decode_config["num-nodes"]:-1}" + decode_config["node-list"]="${decode_config["node-list"]:-}" + decode_config["multiple-workers-per-node"]="${decode_config["multiple-workers-per-node"]:-false}" + decode_config["extra-args"]="${decode_config["extra-args"]:-}" + decode_args["--tensor-parallel-size"]="${decode_args["--tensor-parallel-size"]:-1}" + decode_args["--pipeline-parallel-size"]="${decode_args["--pipeline-parallel-size"]:-1}" + fi +} + _populate_nodelist() { local num_nodes="$1" local exclude_nodelist="$2" @@ -307,7 +348,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 +387,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 @@ -366,12 +418,17 @@ _compute_worker_allocation_vllm() { 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"]}" @@ -421,6 +478,8 @@ _dump_args() { function parse_args() { _parse_cli_pairs "$@" + _normalize_runtime_mode + _init_role_defaults _set_nodelists _patch_dynamo_args @@ -503,7 +562,7 @@ _total_workers_prefill() { } _total_workers_decode() { - echo $(( decode_config["num-nodes"] * decode_config["workers-per-node"] )) + echo $(( ${decode_config["num-nodes"]:-0} * ${decode_config["workers-per-node"]:-0} )) } _count_initialized_prefill() { @@ -511,6 +570,10 @@ _count_initialized_prefill() { } _count_initialized_decode() { + if [[ "${decode_config["num-nodes"]:-0}" -le 0 ]]; then + echo 0 + return + fi grep -i -l -E "${decode_config["worker-initialized-regex"]}" "${RESULTS_DIR}"/dynamo_*decode* 2>/dev/null | wc -l } @@ -571,7 +634,7 @@ _is_frontend_node() { _is_decode_node() { local name="$(_current_node_name)" - [[ ",${decode_config["node-list"]}," == *",$name,"* ]] + [[ ",${decode_config["node-list"]:-}," == *",$name,"* ]] } _is_prefill_node() { @@ -1380,10 +1443,10 @@ function launch_workload() --port "${dynamo_args["port"]}" \ --endpoint "${dynamo_args["endpoint"]}" \ --gpus-per-node "$(_gpus_per_node)" \ - --decode-connector "${decode_args["--connector"]}" \ - --prefill-connector "${prefill_args["--connector"]}" \ + --decode-connector "${decode_args["--connector"]:-}" \ + --prefill-connector "${prefill_args["--connector"]:-}" \ --kvbm-metrics-port "${DYN_KVBM_METRICS_PORT:-6880}" \ - --decode-nodes "${decode_config["node-list"]}" \ + --decode-nodes "${decode_config["node-list"]:-}" \ "${config_arr[@]}" \ -- "${args_arr[@]}" > "${RESULTS_DIR}/$workload_name.log" 2>&1 local workload_status=$? 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 0c09c1f7b..b813cd69f 100644 --- a/src/cloudai/workloads/ai_dynamo/slurm_command_gen_strategy.py +++ b/src/cloudai/workloads/ai_dynamo/slurm_command_gen_strategy.py @@ -406,6 +406,23 @@ def _prepare_aiperf_script(self) -> str | None: script_path.chmod(0o755) return f"{self.CONTAINER_MOUNT_OUTPUT}/{AIPERF_SCRIPT_FILE_NAME}" + def _gen_worker_script_args(self, td: AIDynamoTestDefinition) -> List[str]: + args: List[str] = [] + + 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)}") + + if td.cmd_args.dynamo.mode == "aggregate": + return args + + 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)}") + + return args + def _gen_script_args(self, td: AIDynamoTestDefinition) -> List[str]: self._prepare_hicache_config() self._prepare_lmcache_config() @@ -428,30 +445,21 @@ def _gen_script_args(self, td: AIDynamoTestDefinition) -> List[str]: if td.cmd_args.lmcache_controller: args.append(f"--lmcache-controller-cmd {shlex.quote(td.cmd_args.lmcache_controller.cmd)}") - args.extend( - self._get_toml_args( - td.cmd_args.dynamo, - "--dynamo-", - exclude=[ - "prefill_worker", - "decode_worker", - "dcgm_exporter", - "dcgm-exporter", - ], - ) - ) + dynamo_excluded_fields = [ + "prefill_worker", + "decode_worker", + "dcgm_exporter", + "dcgm-exporter", + ] + if td.cmd_args.dynamo.mode == "disaggregated": + dynamo_excluded_fields.append("mode") + + args.extend(self._get_toml_args(td.cmd_args.dynamo, "--dynamo-", exclude=dynamo_excluded_fields)) if td.cmd_args.dynamo.dcgm_exporter.enabled: args.append('--dynamo-dcgm-exporter-enabled "True"') args.append(f'--dynamo-dcgm-exporter-port "{td.cmd_args.dynamo.dcgm_exporter.port}"') - 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)}") - + args.extend(self._gen_worker_script_args(td)) args.extend(self._get_nested_toml_args(td.cmd_args.genai_perf, "--genai_perf-")) if aiperf_script: args.append(f'--aiperf-name "{td.cmd_args.aiperf.name}"') @@ -656,12 +664,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 = "" 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" diff --git a/tests/workloads/ai_dynamo/test_command_gen_strategy_slurm.py b/tests/workloads/ai_dynamo/test_command_gen_strategy_slurm.py index e0c3d8146..65b2ec9b6 100644 --- a/tests/workloads/ai_dynamo/test_command_gen_strategy_slurm.py +++ b/tests/workloads/ai_dynamo/test_command_gen_strategy_slurm.py @@ -432,6 +432,59 @@ def test_dcgm_exporter_adds_configured_docker_image_installable(cmd_args: AIDyna assert tdef.dcgm_exporter_image in tdef.installables +def test_disaggregate_mode_alias_normalizes() -> None: + args = AIDynamoArgs.model_validate({"mode": "disaggregate"}) + + assert args.mode == "disaggregated" + + +def test_aggregate_mode_uses_combined_vllm_prefill_command() -> None: + args = AIDynamoArgs.model_validate({"mode": "aggregate"}) + + assert args.prefill_worker.cmd == "python3 -m dynamo.vllm" + + +def test_default_disaggregated_mode_preserves_legacy_script_args(strategy: AIDynamoSlurmCommandGenStrategy) -> None: + td = cast(AIDynamoTestDefinition, strategy.test_run.test) + + args = strategy._gen_script_args(td) + command = " ".join(args) + + assert td.cmd_args.dynamo.mode == "disaggregated" + assert '--dynamo-mode "' not in command + assert "--prefill-num-nodes" in command + assert "--prefill-args-model" in command + assert "--decode-num-nodes" in command + assert "--decode-args-model" in command + + +def test_aggregate_mode_omits_decode_script_args(strategy: AIDynamoSlurmCommandGenStrategy) -> None: + td = cast(AIDynamoTestDefinition, strategy.test_run.test) + td.cmd_args.dynamo.mode = "aggregate" + + args = strategy._gen_script_args(td) + command = " ".join(args) + + assert '--dynamo-mode "aggregate"' in args + assert "--prefill-num-nodes" in command + assert "--prefill-args-model" in command + assert "--decode-num-nodes" not in command + assert "--decode-node-list" not in command + assert "--decode-args-" not in command + + +def test_aggregate_mode_uses_prefill_nodes_only( + strategy: AIDynamoSlurmCommandGenStrategy, +) -> None: + td = cast(AIDynamoTestDefinition, strategy.test_run.test) + td.cmd_args.dynamo.mode = "aggregate" + strategy.test_run.nodes = [] + strategy.test_run.num_nodes = 1 + strategy.test_run.num_nodes_explicit = False + + assert strategy.get_cached_nodes_spec()[0] == 1 + + def test_shared_node_disagg_preserves_explicit_smaller_node_count( slurm_system: SlurmSystem, tmp_path: Path, cmd_args: AIDynamoCmdArgs ) -> None: @@ -528,6 +581,17 @@ def test_constraint_allows_shared_node_split_that_fits(slurm_system: SlurmSystem assert td.constraint_check(test_run, slurm_system) +def test_constraint_rejects_aggregate_prefill_that_exceeds_node_gpus( + slurm_system: SlurmSystem, test_run: TestRun +) -> None: + slurm_system.gpus_per_node = 4 + td = cast(AIDynamoTestDefinition, test_run.test) + td.cmd_args.dynamo.mode = "aggregate" + td.cmd_args.dynamo.prefill_worker.args.tensor_parallel_size = 8 + + assert not td.constraint_check(test_run, slurm_system) + + def test_constraint_rejects_shared_node_split_that_exceeds_node_gpus( slurm_system: SlurmSystem, test_run: TestRun ) -> None: