Merge branch 'main' into ishan/unify-trace

Signed-off-by: ishandhanani <[email protected]>

Author: ishandhanani

.github/actions/docker-build/action.yml

@@ -151,6 +151,26 @@ runs:
         # Exit with the build's exit code
         exit ${BUILD_EXIT_CODE}
 
+    - name: Run Sanity Check on Runtime Image
+      if: inputs.target == 'runtime'
+      shell: bash
+      run: |
+        IMAGE_TAG="${{ steps.build.outputs.image_tag }}"
+        echo "Running sanity check on image: $IMAGE_TAG"
+
+        # Run the sanity check script inside the container
+        # The script is located in /workspace/deploy/sanity_check.py in runtime containers
+        set +e
+        docker run --rm "$IMAGE_TAG" python /workspace/deploy/sanity_check.py --runtime-check --no-gpu-check
+        SANITY_CHECK_EXIT_CODE=$?
+        set -e
+        if [ ${SANITY_CHECK_EXIT_CODE} -ne 0 ]; then
+          echo "ERROR: Sanity check failed - ai-dynamo packages not properly installed"
+          exit ${SANITY_CHECK_EXIT_CODE}
+        else
+          echo "✅ Sanity check passed"
+        fi
+
     - name: Capture Build Metrics
       id: metrics
       shell: bash

components/src/dynamo/planner/utils/planner_core.py

@@ -24,7 +24,7 @@
     PrefillInterpolator,
 )
 from dynamo.planner.utils.pre_swept_results_utils import PreSweptResultsHelper
-from dynamo.planner.utils.prometheus import PrometheusAPIClient
+from dynamo.planner.utils.prometheus import MetricSource, PrometheusAPIClient
 from dynamo.planner.utils.trace_data_extractor import extract_metrics_from_mooncake
 from dynamo.runtime import DistributedRuntime
 from dynamo.runtime.logging import configure_dynamo_logging
@@ -150,9 +150,20 @@ def __init__(
                 else:
                     raise ValueError(f"Invalid environment: {args.environment}")
 
+            # Use backend metrics for vLLM (queries vllm:* metrics directly from workers)
+            # Use frontend metrics for other backends (queries dynamo_frontend_* metrics)
+            metric_source = (
+                MetricSource.VLLM
+                if args.backend.lower() == "vllm"
+                else MetricSource.FRONTEND
+            )
+            logger.info(
+                f"Initializing Prometheus client with metric_source='{metric_source}' for backend '{args.backend}'"
+            )
             self.prometheus_api_client = PrometheusAPIClient(
                 args.metric_pulling_prometheus_endpoint,
                 args.namespace,
+                metric_source=metric_source,
             )
 
         self.num_req_predictor = LOAD_PREDICTORS[args.load_predictor](