add remaning fork functionality

Author: podkidyshev

conf/experimental/ai_dynamo/test/vllm.toml

@@ -74,11 +74,16 @@ workloads = "aiperf.sh"
     concurrency = 2
 
   [cmd_args.aiperf]
+  health-check-between-phases = true
+  continue-on-phase-failure = false
     [cmd_args.aiperf.args]
     concurrency = 2
+    endpoint-type = "chat"
     extra-inputs = '{"min_tokens":10}'
     output-tokens-mean = 500
     request-count = 50
+    server-metrics = "auto"
+    streaming = true
     synthetic-input-tokens-mean = 300
 
   [[cmd_args.aiperf_phases]]
@@ -94,6 +99,7 @@ workloads = "aiperf.sh"
     [cmd_args.aiperf_phases.args]
     concurrency = 4
     request-count = 50
+    streaming = false
 
   [cmd_args.aiperf_accuracy]
   entrypoint = "aiperf profile"

conf/experimental/ai_dynamo/test_scenario/vllm_lmcache.toml

@@ -24,7 +24,10 @@ description = "Self-contained AIDynamo scenario wiring vLLM disaggregated infere
 test_template_name = "AIDynamo"
 time_limit = "00:10:00"
 extra_container_mounts = ["/run/udev:/run/udev"]
-dse_excluded_args = ["cmd_args.lmcache.lmcache_worker_ports", "cmd_args.aiperf_phases"]
+dse_excluded_args = [
+  "cmd_args.lmcache.lmcache_worker_ports",
+  "cmd_args.aiperf_phases",
+]
 
   [Tests.cmd_args]
   docker_image_url = "nvcr.io/nvidia/ai-dynamo/vllm-runtime:1.1.1"

conf/experimental/ai_dynamo/test_scenario/vllm_slurm.toml

@@ -36,12 +36,28 @@ time_limit = "00:10:00"
       tensor-parallel-size = 4
       pipeline-parallel-size = 1
 
+    [[Tests.cmd_args.aiperf_phases]]
+    name = "round_1"
+      [Tests.cmd_args.aiperf_phases.args]
+      concurrency = 2
+      request-count = 50
+      server-metrics = "auto"
+
+    [[Tests.cmd_args.aiperf_phases]]
+    name = "round_2"
+      [Tests.cmd_args.aiperf_phases.args]
+      concurrency = 4
+      request-count = 50
+
 [[Tests]]
 id = "test.disagg.multinode"
 test_name = "vLLM"
 time_limit = "00:10:00"
 
   [Tests.cmd_args]
+    [Tests.cmd_args.dynamo.dcgm_exporter]
+    enabled = true
+
     [Tests.cmd_args.dynamo.prefill_worker]
     num-nodes = 2
       [Tests.cmd_args.dynamo.prefill_worker.args]
@@ -53,3 +69,17 @@ time_limit = "00:10:00"
       [Tests.cmd_args.dynamo.decode_worker.args]
       tensor-parallel-size = 4
       pipeline-parallel-size = 1
+
+    [[Tests.cmd_args.aiperf_phases]]
+    name = "round_1"
+      [Tests.cmd_args.aiperf_phases.args]
+      concurrency = 4
+      request-count = 50
+      server-metrics = "auto"
+
+    [[Tests.cmd_args.aiperf_phases]]
+    name = "round_2"
+      [Tests.cmd_args.aiperf_phases.args]
+      concurrency = 8
+      request-count = 50
+      server-metrics = "auto"

src/cloudai/workloads/ai_dynamo/__init__.py

@@ -23,6 +23,7 @@
     AIPerf,
     AIPerfAccuracy,
     AIPerfPhase,
+    DCGMExporter,
     GenAIPerf,
     LMCacheController,
     WorkerBaseArgs,
@@ -44,6 +45,7 @@
     "AIPerf",
     "AIPerfAccuracy",
     "AIPerfPhase",
+    "DCGMExporter",
     "GenAIPerf",
     "LMCacheController",
     "WorkerBaseArgs",

src/cloudai/workloads/ai_dynamo/ai_dynamo.py

@@ -140,6 +140,20 @@ class WorkerConfig(BaseModel):
     )
 
 
+class DCGMExporter(BaseModel):
+    """Optional DCGM exporter launch configuration."""
+
+    model_config = ConfigDict(extra="forbid", populate_by_name=True)
+
+    enabled: bool = False
+    image_url: str = Field(
+        default="nvcr.io/nvidia/k8s/dcgm-exporter:4.5.2-4.8.1-distroless",
+        serialization_alias="image-url",
+        validation_alias=AliasChoices("image-url", "image_url"),
+    )
+    port: int = 9401
+
+
 class AIDynamoArgs(BaseModel):
     """Arguments for AI Dynamo setup."""
 
@@ -205,6 +219,7 @@ def validate_connector(cls, v: str | list[str] | None) -> str | list[str] | None
         serialization_alias="nats-port",
         validation_alias=AliasChoices("nats-port", "nats_port"),
     )
+    dcgm_exporter: DCGMExporter = Field(default_factory=DCGMExporter)
 
     decode_worker: WorkerConfig = WorkerConfig(
         cmd="python3 -m dynamo.vllm",
@@ -269,6 +284,16 @@ class AIPerf(Workload):
         serialization_alias="artifact-dir-name",
         validation_alias=AliasChoices("artifact-dir-name", "artifact_dir_name"),
     )
+    health_check_between_phases: bool = Field(
+        default=True,
+        serialization_alias="health-check-between-phases",
+        validation_alias=AliasChoices("health-check-between-phases", "health_check_between_phases"),
+    )
+    continue_on_phase_failure: bool = Field(
+        default=False,
+        serialization_alias="continue-on-phase-failure",
+        validation_alias=AliasChoices("continue-on-phase-failure", "continue_on_phase_failure"),
+    )
 
     @property
     def installables(self) -> list[Installable]:

src/cloudai/workloads/ai_dynamo/ai_dynamo.sh

@@ -60,6 +60,8 @@ dynamo_args["worker-error-pattern"]="zmq.error.ZMQError:.Address.already.in.use|
 dynamo_args["sgl-http-port"]=9001
 dynamo_args["prefill-port"]=30011
 dynamo_args["decode-port"]=30021
+dynamo_args["dcgm-exporter-enabled"]="False"
+dynamo_args["dcgm-exporter-port"]=9401
 
 function log()
 {
@@ -892,6 +894,38 @@ _query_frontend() {
   curl -s -X POST "${dynamo_args["url"]}/v1/chat/completions" -H "Content-Type: application/json" -d @$RESULTS_DIR/curl_cmd.json
 }
 
+_resolve_aiperf_server_metrics_urls() {
+  local urls="http://${dynamo_args["frontend-node"]}:${dynamo_args["port"]}/metrics"
+  local base_system_port=${DYN_SYSTEM_PORT:-9090}
+  local decode_workers_per_node=${decode_config["workers-per-node"]:-1}
+  local prefill_workers_per_node=${prefill_config["workers-per-node"]:-1}
+  local IFS_SAVE="$IFS"
+  local node i
+
+  IFS=','
+  for node in ${prefill_config["node-list"]:-}; do
+    for i in $(seq 0 $(( prefill_workers_per_node - 1 ))); do
+      urls="${urls},http://${node}:$((base_system_port + i))/metrics"
+    done
+  done
+
+  for node in ${decode_config["node-list"]:-}; do
+    for i in $(seq 0 $(( decode_workers_per_node - 1 ))); do
+      urls="${urls},http://${node}:$((base_system_port + i))/metrics"
+    done
+  done
+
+  if [[ "${dynamo_args["dcgm-exporter-enabled"]}" == "True" || "${dynamo_args["dcgm-exporter-enabled"]}" == "true" ]]; then
+    for node in ${decode_config["node-list"]:-},${prefill_config["node-list"]:-}; do
+      [[ -z "$node" ]] && continue
+      urls="${urls},http://${node}:${dynamo_args["dcgm-exporter-port"]}/metrics"
+    done
+  fi
+  IFS="$IFS_SAVE"
+
+  echo "$urls"
+}
+
 function setup_cufile()
 {
   export CUFILE_ENV_PATH_JSON="$RESULTS_DIR/cufile.json"
@@ -1059,6 +1093,10 @@ function launch_workload()
   local workload_name="${workload_config_ref["--name"]}"
   local script="${workload_config_ref["--script"]}"
   export FRONTEND_URL="${dynamo_args["url"]}"
+  export AIPERF_MODEL="${dynamo_args["model"]}"
+  export AIPERF_ENDPOINT="${dynamo_args["endpoint"]}"
+  export AIPERF_FAILURE_MARKER="${FATAL_ERROR_MARKER}"
+  export AIPERF_SERVER_METRICS_URLS="$(_resolve_aiperf_server_metrics_urls)"
 
   # Build config and workload args as proper bash arrays to preserve
   # multi-word values (e.g. --cmd "genai-perf profile") through word splitting.