Merge #134390

134390: drtprod: modify tpcc run script. r=vidit-bhat a=shailendra-patel Modifying generate_tpcc_run.sh script to equally distribute pgurls of cluster nodes among workload nodes. Epic: none Release note: None Co-authored-by: Shailendra Patel <[email protected]>
cockroachdb · Nov 21, 2024 · feec8b9 · feec8b9
2 parents 6a7b1e2 + d07f09f
commit feec8b9
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 16 deletions.
diff --git a/pkg/cmd/drtprod/configs/drt_scale.yaml b/pkg/cmd/drtprod/configs/drt_scale.yaml
@@ -162,8 +162,8 @@ targets:
         flags:
           db: cct_tpcc
           warehouses: 400000
-          max-rate: 10000
-          workers: 5000
+          max-rate: 1000
+          workers: 400000
           conns: 5000
           duration: 12h
           ramp: 10m

diff --git a/pkg/cmd/drtprod/scripts/generate_tpcc_run.sh b/pkg/cmd/drtprod/scripts/generate_tpcc_run.sh
@@ -39,20 +39,32 @@ if [ -z "${WORKLOAD_NODES}" ]; then
   exit 1
 fi
 
+if [ -z "${CLUSTER_NODES}" ]; then
+  echo "environment CLUSTER_NODES is not set"
+  exit 1
+fi
+
 absolute_path=$(roachprod run "${WORKLOAD_CLUSTER}":1 -- "realpath ./cockroach")
 pwd=$(roachprod run "${WORKLOAD_CLUSTER}":1 -- "dirname ${absolute_path}")
-# Prepare PGURLS
-PGURLS=$(roachprod pgurl $CLUSTER | sed s/\'//g)
 
-# Loop through each node
-for NODE in $(seq 1 $WORKLOAD_NODES)
-do
+# Calculate the number of PGURLS each workload node should get
+PGURL_PER_NODE=$((CLUSTER_NODES / WORKLOAD_NODES))
+REMAINDER_NODE=$((CLUSTER_NODES % WORKLOAD_NODES))
+
+# Distribute the PGURLS among the workload nodes
+for ((NODE=0; NODE<WORKLOAD_NODES; NODE++)); do
+  START_OFFSET=$((NODE * PGURL_PER_NODE + (NODE < REMAINDER_NODE ? NODE : REMAINDER_NODE) + 1))
+  END_OFFSET=$((START_OFFSET + PGURL_PER_NODE + (NODE < REMAINDER_NODE ? 1 : 0) - 1))
+
+  # Print or use the PGURLS for the current workload node
+  echo "pgurl for Nodes ${START_OFFSET}:${END_OFFSET}"
+  PGURLS=$(roachprod pgurl $CLUSTER:$START_OFFSET-$END_OFFSET)
+
+
   # Create the workload script
   cat <<EOF >/tmp/tpcc_run_${suffix}.sh
 #!/usr/bin/env bash
 
-read -r -a PGURLS_ARR <<< "$PGURLS"
-
 j=0
 while true; do
     echo ">> Starting tpcc workload"
@@ -61,20 +73,20 @@ while true; do
     ./cockroach workload run tpcc $@ \
         --tolerate-errors \
         --families \
-        "\${PGURLS_ARR[@]}" | tee \$LOG
+         $PGURLS  | tee \$LOG
     if [ \$? -eq 0 ]; then
         rm "\$LOG"
     fi
     sleep 1
 done
 EOF
 
-  # Upload the script to the workload cluster
-  roachprod put $WORKLOAD_CLUSTER:$NODE /tmp/tpcc_run_${suffix}.sh
-  roachprod ssh $WORKLOAD_CLUSTER:$NODE -- "chmod +x tpcc_run_${suffix}.sh"
+#   Upload the script to the workload cluster
+  roachprod put $WORKLOAD_CLUSTER:$((NODE + 1)) /tmp/tpcc_run_${suffix}.sh
+  roachprod ssh $WORKLOAD_CLUSTER:$((NODE + 1)) -- "chmod +x tpcc_run_${suffix}.sh"
   if [ "$execute_script" = "true" ]; then
-    roachprod run "${WORKLOAD_CLUSTER}":1 -- "sudo systemd-run --unit tpcc_run_${suffix} --same-dir --uid \$(id -u) --gid \$(id -g) bash ${pwd}/tpcc_run_${suffix}.sh"
+      roachprod run "${WORKLOAD_CLUSTER}":1 -- "sudo systemd-run --unit tpcc_run_${suffix} --same-dir --uid \$(id -u) --gid \$(id -g) bash ${pwd}/tpcc_run_${suffix}.sh"
   else
-    echo "Run --> roachprod run "${WORKLOAD_CLUSTER}":1 -- \"sudo systemd-run --unit tpcc_run_${suffix} --same-dir --uid \\\$(id -u) --gid \\\$(id -g) bash ${pwd}/tpcc_run_${suffix}.sh\""
-  fi
+    echo "Run --> roachprod run "${WORKLOAD_CLUSTER}":$((NODE + 1)) -- \"sudo systemd-run --unit tpcc_run_${suffix} --same-dir --uid \\\$(id -u) --gid \\\$(id -g) bash ${pwd}/tpcc_run_${suffix}.sh\""
+    fi
 done