Feature/DCGM Exporter (#98)

Author: burakemirsezen

src/swiss_ai_model_launch/assets/template.jinja

@@ -27,6 +27,10 @@ OCF_BOOTSTRAP_ADDR="/ip4/148.187.108.178/tcp/43905/p2p/QmbUKJkCfotDzbFE5uoTsXD4G
 SERVED_MODEL_NAME="{{ served_model_name }}"
 METRICS_REMOTE_WRITE_URL="{{ metrics_remote_write_url or '' }}"
 METRICS_AGENT_BIN="{{ metrics_agent_binary }}"
+DCGM_EXPORTER_BIN="{{ dcgm_exporter_binary }}"
+USE_DCGM_EXPORTER={{ "false" if disable_dcgm_exporter else "true" }}
+USE_METRICS={{ "false" if disable_metrics else "true" }}
+
 
 {% if telemetry_endpoint %}
 curl -sf -X POST "{{ telemetry_endpoint }}" \
@@ -44,11 +48,13 @@ if [[ "$ARCH" == "aarch64" ]]; then
     export SP_NCCL_SO_PATH=/usr/lib/aarch64-linux-gnu/
     export OCF_BIN=/ocfbin/ocf-arm
     METRICS_AGENT_BIN="${METRICS_AGENT_BIN}-arm64"
+    DCGM_EXPORTER_BIN="${DCGM_EXPORTER_BIN}-arm64"
 elif [[ "$ARCH" == "x86_64" ]]; then
     echo "Running on x86_64"
     export SP_NCCL_SO_PATH=/usr/lib/x86_64-linux-gnu/
     export OCF_BIN=/ocfbin/ocf-amd64
     METRICS_AGENT_BIN="${METRICS_AGENT_BIN}-amd64"
+    DCGM_EXPORTER_BIN="${DCGM_EXPORTER_BIN}-amd64"
 else
     echo "Unknown architecture: $ARCH"
     exit 1
@@ -197,23 +203,69 @@ $FRAMEWORK_CMD" &
     done
 done
 
-# Push framework metrics to Prometheus via vmagent (runs on the batch node).
-# Pyxis containers share the host network namespace, so localhost:8080 on the
-# batch node reaches the framework's API server. The scrape config is a static
-# YAML staged alongside the binaries at /ocfbin/vmagent-scrape.yaml.
-# NOTE: only the worker on the batch node is scraped — multi-worker setups
-# would need a per-worker vmagent (future work).
+# Launch vmagent and DCGM on every node.
+# The batch node (index 0) runs directly; worker nodes run via srun --overlap.
+# The batch node scrapes framework metrics (8080) + DCGM (9400); workers scrape only DCGM (9400).
 if [ -n "$METRICS_REMOTE_WRITE_URL" ] && [ -x "$METRICS_AGENT_BIN" ]; then
-    "$METRICS_AGENT_BIN" \
-        -promscrape.config=/capstor/store/cscs/swissai/infra01/ocf-share/vmagent-scrape.yaml \
-        -remoteWrite.url="${METRICS_REMOTE_WRITE_URL}" \
-        -remoteWrite.label="slurm_job_id=${SLURM_JOB_ID}" \
-        -remoteWrite.label="model=${SERVED_MODEL_NAME}" \
-        -remoteWrite.label="framework=${FRAMEWORK}" \
-        -remoteWrite.label="user=${USER}" \
-        -remoteWrite.tmpDataPath=/tmp/vmagent-data-${SLURM_JOB_ID} \
-        > /tmp/vmagent-${SLURM_JOB_ID}.log 2>&1 &
-elif [ -n "$METRICS_REMOTE_WRITE_URL" ]; then
+    VMAGENT_COMMON_ARGS="
+        -remoteWrite.url=${METRICS_REMOTE_WRITE_URL}
+        -remoteWrite.label=slurm_job_id=${SLURM_JOB_ID}
+        -remoteWrite.label=model=${SERVED_MODEL_NAME}
+        -remoteWrite.label=framework=${FRAMEWORK}
+        -remoteWrite.label=user=${USER}
+    "
+    METRICS_CONFIG_DIR="/capstor/store/cscs/swissai/infra01/ocf-share"
+    DCGM_COMMON_ARGS="--address 0.0.0.0:9400 -f $METRICS_CONFIG_DIR/default-counters.csv"
+    DCGM_LOG="/tmp/dcgm-exporter-${SLURM_JOB_ID}.log"
+    VMAGENT_LOG="/tmp/vmagent-${SLURM_JOB_ID}.log"
+    VMAGENT_DATA="/tmp/vmagent-data-${SLURM_JOB_ID}"
+
+    if [ "$USE_METRICS" != "true" ]; then
+        echo "Metrics disabled, skipping vmagent initialization" >&2
+    else
+        for i in "${!nodes[@]}"; do
+            node="${nodes[$i]}"
+            if [ "$i" -eq 0 ]; then
+                if [ "$USE_DCGM_EXPORTER" = "true" ]; then
+                    echo "dcgm-exporter: /dev/nvidia0 exists=$([ -e /dev/nvidia0 ] && echo yes || echo no), binary=$DCGM_EXPORTER_BIN executable=$([ -x "$DCGM_EXPORTER_BIN" ] && echo yes || echo no)" >&2
+                    if [ -e /dev/nvidia0 ] && [ -x "$DCGM_EXPORTER_BIN" ]; then
+                        "$DCGM_EXPORTER_BIN" $DCGM_COMMON_ARGS > $DCGM_LOG 2>&1 &
+                    else
+                        echo "dcgm-exporter: no NVIDIA GPU or binary not found, skipping" >&2
+                    fi
+                    VMAGENT_SCRAPE_CONFIG="$METRICS_CONFIG_DIR/vmagent-scrape.yaml"
+                else
+                    echo "dcgm-exporter: disabled (USE_DCGM_EXPORTER != true), skipping" >&2
+                    VMAGENT_SCRAPE_CONFIG="$METRICS_CONFIG_DIR/vmagent-scrape-no-dcgm.yaml"
+                fi
+                "$METRICS_AGENT_BIN" $VMAGENT_COMMON_ARGS \
+                    -promscrape.config=$VMAGENT_SCRAPE_CONFIG \
+                    -remoteWrite.label="node=$(hostname)" \
+                    -remoteWrite.tmpDataPath=$VMAGENT_DATA \
+                    > $VMAGENT_LOG 2>&1 &
+            else
+                srun --nodes=1 --ntasks=1 --nodelist=$node --overlap \
+                    bash -c "
+                        if [ \"$USE_DCGM_EXPORTER\" = 'true' ]; then
+                            if [ -e /dev/nvidia0 ] && [ -x \"$DCGM_EXPORTER_BIN\" ]; then
+                                \"$DCGM_EXPORTER_BIN\" $DCGM_COMMON_ARGS > $DCGM_LOG 2>&1 &
+                            else
+                                echo 'dcgm-exporter: no NVIDIA GPU or binary not found, skipping' >&2
+                            fi
+                            \"$METRICS_AGENT_BIN\" $VMAGENT_COMMON_ARGS \
+                                -promscrape.config=$METRICS_CONFIG_DIR/vmagent-scrape-dcgm-only.yaml \
+                                -remoteWrite.label=\"node=\$(hostname)\" \
+                                -remoteWrite.tmpDataPath=$VMAGENT_DATA \
+                                > $VMAGENT_LOG 2>&1 &
+                        else
+                            echo 'dcgm-exporter: disabled (USE_DCGM_EXPORTER != true), skipping vmagent on worker node' >&2
+                        fi
+                        wait
+                    " &
+            fi
+        done
+    fi
+else
     echo "metrics: $METRICS_AGENT_BIN not found, skipping push" >&2
 fi
 

src/swiss_ai_model_launch/cli/main.py

@@ -246,6 +246,18 @@ def _build_parser() -> argparse.ArgumentParser:
         action="store_true",
         help="Disable OCF.",
     )
+    advanced_parser.add_argument(
+        "--disable-dcgm-exporter",
+        dest="disable_dcgm_exporter",
+        action="store_true",
+        help="Disable the DCGM exporter.",
+    )
+    advanced_parser.add_argument(
+        "--disable-metrics",
+        dest="disable_metrics",
+        action="store_true",
+        help="Disable metrics collection.",
+    )
     advanced_parser.add_argument(
         "--pre-launch-cmds",
         dest="pre_launch_cmds",
@@ -550,6 +562,8 @@ async def _run_advanced(args: argparse.Namespace) -> None:
         use_router=args.use_router,
         router_args=args.router_args,
         disable_ocf=args.disable_ocf,
+        disable_dcgm_exporter=args.disable_dcgm_exporter,
+        disable_metrics=args.disable_metrics,
         telemetry_endpoint=config.get_value("telemetry_endpoint"),
     )
 

src/swiss_ai_model_launch/launchers/launch_args.py

@@ -25,6 +25,9 @@ class LaunchArgs(BaseModel):
     telemetry_endpoint: str | None = None
     metrics_remote_write_url: str = "https://prometheus-dev.swissai.svc.cscs.ch/api/v1/write"
     metrics_agent_binary: str = "/capstor/store/cscs/swissai/infra01/ocf-share/vmagent"
+    dcgm_exporter_binary: str = "/capstor/store/cscs/swissai/infra01/ocf-share/dcgm-exporter"
+    disable_dcgm_exporter: bool = False
+    disable_metrics: bool = False
 
     @model_validator(mode="after")
     def set_defaults(self) -> "LaunchArgs":