From 3cc5171fba910250bed990ba4fe305557061ff2c Mon Sep 17 00:00:00 2001
From: Paulo Castello da Costa <pcastello@google.com>
Date: Mon, 31 Jul 2023 10:49:44 -0700
Subject: [PATCH] [Test] Write Prometheus stats to driver log. (#367)

Write Prometheus stats to driver log, in addition to worker stats, even
when not saving data to BigQuery.
---
 containers/runtime/driver/Dockerfile |  3 +++
 containers/runtime/driver/run.sh     | 20 +++++++++++---------
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/containers/runtime/driver/Dockerfile b/containers/runtime/driver/Dockerfile
index 110f0407..15758b74 100644
--- a/containers/runtime/driver/Dockerfile
+++ b/containers/runtime/driver/Dockerfile
@@ -93,6 +93,9 @@ RUN pip3 install \
 COPY . /src/driver
 RUN chmod a+x /src/driver/run.sh
 
+# TODO(paulosjca): Remove this when prometheus.py is executable at driver commit.
+RUN chmod a+x /src/code/tools/run_tests/performance/prometheus.py
+
 ENV QPS_WORKERS=""
 ENV QPS_WORKERS_FILE=""
 ENV SCENARIOS_FILE="/src/driver/example.json"
diff --git a/containers/runtime/driver/run.sh b/containers/runtime/driver/run.sh
index 20726d7b..581b49ad 100644
--- a/containers/runtime/driver/run.sh
+++ b/containers/runtime/driver/run.sh
@@ -29,20 +29,22 @@ fi
 
 /src/code/bazel-bin/test/cpp/qps/qps_json_driver --quit=true
 
+if [ -n "${SERVER_TARGET_OVERRIDE}" ] || [ -n "${ENABLE_PROMETHEUS}" ]; then
+  if  [ "$(dig +short -t srv prometheus.test-infra-system.svc.cluster.local)" ]; then
+    /src/code/tools/run_tests/performance/prometheus.py \
+      --url=http://prometheus.test-infra-system.svc.cluster.local:9090 \
+      --pod_type=clients --container_name=main \
+      --container_name=sidecar --delay_seconds=20 \
+      --export_file_name=prometheus_query_results.json
+  fi
+fi
+
 if [ -n "${BQ_RESULT_TABLE}" ]; then
   if [ -r "${METADATA_OUTPUT_FILE}" ]; then
     cp "${METADATA_OUTPUT_FILE}" metadata.json
   fi
   if [ -r "${NODE_INFO_OUTPUT_FILE}" ]; then
     cp "${NODE_INFO_OUTPUT_FILE}" node_info.json
-    if [ -n "${SERVER_TARGET_OVERRIDE}" ] || [ -n "${ENABLE_PROMETHEUS}" ]; then
-      if  [ "$(dig +short -t srv prometheus.test-infra-system.svc.cluster.local)" ]; then
-        python3 /src/code/tools/run_tests/performance/prometheus.py \
-          --url=http://prometheus.test-infra-system.svc.cluster.local:9090 \
-          --pod_type=clients --container_name=main \
-          --container_name=sidecar --delay_seconds=20
-      fi
-    fi
   fi
-  python3 /src/code/tools/run_tests/performance/bq_upload_result.py --bq_result_table="${BQ_RESULT_TABLE}"
+  /src/code/tools/run_tests/performance/bq_upload_result.py --bq_result_table="${BQ_RESULT_TABLE}"
 fi