Misiug/slurmprofiles

.gitignore

@@ -22,4 +22,14 @@ presto/docker/config/generated*/
 # Generated Presto Docker Compose files
 presto/docker/docker-compose/generated*/
 
+# Slurm logs and results
+presto/slurm/presto-nvl72/logs/
+presto/slurm/presto-nvl72/*.err
+presto/slurm/presto-nvl72/*.out
+presto/slurm/presto-nvl72/result_dir/
+presto/slurm/presto-nvl72/kept_results/
+presto/slurm/presto-nvl72/worker_data/
+presto/slurm/presto-nvl72/profiles/
+presto/slurm/presto-nvl72/worker_info/
+
 devstate*

benchmark_data_tools/duckdb_utils.py

@@ -6,6 +6,10 @@
 import duckdb
 
 
+def quote_ident(name: str) -> str:
+    return '"' + name.replace('"', '""') + '"'
+
+
 def init_benchmark_tables(benchmark_type, scale_factor):
     tables = duckdb.sql("SHOW TABLES").fetchall()
     assert len(tables) == 0
@@ -22,27 +26,27 @@ def init_benchmark_tables(benchmark_type, scale_factor):
 def drop_benchmark_tables():
     tables = duckdb.sql("SHOW TABLES").fetchall()
     for (table,) in tables:
-        duckdb.sql(f"DROP TABLE {table}")
+        duckdb.sql(f"DROP TABLE {quote_ident(table)}")
 
 
 def create_table(table_name, data_path):
-    duckdb.sql(f"DROP TABLE IF EXISTS {table_name}")
-    duckdb.sql(f"CREATE TABLE {table_name} AS SELECT * FROM '{data_path}/*.parquet';")
+    duckdb.sql(f"DROP TABLE IF EXISTS {quote_ident(table_name)}")
+    duckdb.sql(f"CREATE TABLE {quote_ident(table_name)} AS SELECT * FROM '{data_path}/*.parquet';")
 
 
 # Generates a sample table with a small limit.
 # This is mainly used to extract the schema from the parquet files.
 def create_not_null_table_from_sample(table_name, data_path):
-    duckdb.sql(f"DROP TABLE IF EXISTS {table_name}")
-    duckdb.sql(f"CREATE TABLE {table_name} AS SELECT * FROM '{data_path}/*.parquet' LIMIT 10;")
-    ret = duckdb.sql(f"DESCRIBE TABLE {table_name}").fetchall()
+    duckdb.sql(f"DROP TABLE IF EXISTS {quote_ident(table_name)}")
+    duckdb.sql(f"CREATE TABLE {quote_ident(table_name)} AS SELECT * FROM '{data_path}/*.parquet' LIMIT 10;")
+    ret = duckdb.sql(f"DESCRIBE TABLE {quote_ident(table_name)}").fetchall()
     for row in ret:
-        duckdb.sql(f"ALTER TABLE {table_name} ALTER COLUMN {row[0]} SET NOT NULL;")
+        duckdb.sql(f"ALTER TABLE {quote_ident(table_name)} ALTER COLUMN {row[0]} SET NOT NULL;")
 
 
 def create_table_from_sample(table_name, data_path):
-    duckdb.sql(f"DROP TABLE IF EXISTS {table_name}")
-    duckdb.sql(f"CREATE TABLE {table_name} AS SELECT * FROM '{data_path}/*.parquet' LIMIT 10;")
+    duckdb.sql(f"DROP TABLE IF EXISTS {quote_ident(table_name)}")
+    duckdb.sql(f"CREATE TABLE {quote_ident(table_name)} AS SELECT * FROM '{data_path}/*.parquet' LIMIT 10;")
 
 
 def is_decimal_column(column_type):

presto/docker/config/template/etc_coordinator/config_native.properties

@@ -48,7 +48,7 @@ optimizer.generate-domain-filters=true
 # Upper limit for broadcasted table size to avoid memory blowups.
 # See: https://github.com/prestodb/presto/issues/22161#issuecomment-1994128619
 join-max-broadcast-table-size={{ .JoinMaxBroadcastTableSizeMb }}MB
-# Default is AUTOMATIC, ucx exchange does not support BROADCAST partition type.
+# overwritten to "PARTITIONED" in multi-node context
 join-distribution-type=AUTOMATIC
 
 # Client request timeout to avoid hung queries.
@@ -58,7 +58,7 @@ query.execution-policy=phased
 # Kill queries based on total reservation on blocked nodes to recover memory.
 query.low-memory-killer.policy=total-reservation-on-blocked-nodes
 # Upper limit on query wall time to keep tests bounded.
-query.max-execution-time=30m
+query.max-execution-time=10m
 # Keep metadata of up to 1000 queries for UI and debugging.
 query.max-history=1000
 # Memory quotas per node and cluster to protect stability.

presto/docker/config/template/etc_worker/config_native.properties

@@ -29,13 +29,15 @@ system-mem-limit-gb={{ sub .ContainerMemoryGb .GeneratorParameters.MemoryPushBac
 system-mem-shrink-gb=20
 
 # Optimize for single-node execution when the entire query can run locally.
+# overwritten to "false" in multi-node settings.
 single-node-execution-enabled=true
 
 # Enable cuDF (CPU mode will ignore this setting)
 cudf.enabled=true
+# overwritten to "true" in multi-node settings.
 cudf.exchange=false
-# Port number currently must be exactly 3 more than server port (ignored if cudf.exchange is false)
-cudf.exchange.server.port=8083
+# overwritten when cudf.exchange is enabled (ignored otherwise)
+cudf.exchange.server.port=0000
 cudf.memory_resource=async
 
 async-data-cache-enabled=false

presto/scripts/common_functions.sh

@@ -1,25 +1,12 @@
 #!/bin/bash
 
-# Copyright (c) 2025, NVIDIA CORPORATION.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+# SPDX-License-Identifier: Apache-2.0
 
 function wait_for_worker_node_registration() {
   trap "rm -rf node_response.json" RETURN
 
   echo "Waiting for a worker node to be registered..."
-  HOSTNAME=${1:-localhost}
-  PORT=${2:-8080}
   COORDINATOR_URL=http://${HOSTNAME}:${PORT}
   echo "Coordinator URL: $COORDINATOR_URL"
   local -r MAX_RETRIES=12

presto/scripts/generate_presto_config.sh

@@ -30,19 +30,23 @@ if [ ! -x "${SCRIPT_DIR}/../pbench/pbench" ]; then
   echo_error "ERROR: generate_presto_config.sh script cannot find pbench at ${SCRIPT_DIR}/../pbench/pbench"
 fi
 
+# This function duplicates the worker configs when we are running multiple workers.
+# It also adds certain config options to the workers if those options apply only to multi-worker environments.
 function duplicate_worker_configs() {
   local worker_id=$1
   echo "Duplicating worker configs for GPU ID $worker_id"
   local worker_config="${CONFIG_DIR}/etc_worker_${worker_id}"
-  local worker_native_config="${worker_config}/config_native.properties"
   local coord_config="${CONFIG_DIR}/etc_coordinator"
+  local worker_native_config="${worker_config}/config_native.properties"
   local coord_native_config="${coord_config}/config_native.properties"
+  # Need to stagger the port numbers because ucx exchange currently expects to be exactly
+  # 3 higher than the http port.
   local http_port="10$(printf "%02d\n" "$worker_id")0"
   local exch_port="10$(printf "%02d\n" "$worker_id")3"
   rm -rf ${worker_config}
   cp -r ${CONFIG_DIR}/etc_worker ${worker_config}
 
-  # Single node execution needs to be disabled if we are running multiple workers.
+  # Some configs should only be applied if we are in a multi-worker environment.
   if [[ ${NUM_WORKERS} -gt 1 ]]; then
     sed -i "s+single-node-execution-enabled.*+single-node-execution-enabled=false+g" ${coord_native_config}
     sed -i "s+single-node-execution-enabled.*+single-node-execution-enabled=false+g" ${worker_native_config}
@@ -72,7 +76,7 @@ RAM_GB=$(lsmem -b | grep "Total online memory" | awk '{print int($4 / (1024*1024
 if [[ -z ${VARIANT_TYPE} || ! ${VARIANT_TYPE} =~ ^(cpu|gpu|java)$ ]]; then
   echo_error "ERROR: VARIANT_TYPE must be set to a valid variant type (cpu, gpu, java)."
 fi
-if [[ -z ${VCPU_PER_WORKER} ]]; then
+if [[ -z ${VCPU_PER_WORKER:-} ]]; then
   if [[ "${VARIANT_TYPE}" == "gpu" ]]; then
     VCPU_PER_WORKER=2
   else
@@ -122,6 +126,7 @@ EOF
   fi
 
   COORD_CONFIG="${CONFIG_DIR}/etc_coordinator/config_native.properties"
+  WORKER_CONFIG="${CONFIG_DIR}/etc_worker/config_native.properties"
   # now perform other variant-specific modifications to the generated configs
   if [[ "${VARIANT_TYPE}" == "gpu" ]]; then
     # for GPU variant, uncomment these optimizer settings
@@ -158,10 +163,13 @@ fi
 
 # We want to propagate any changes from the original worker config to the new worker configs even if
 # we did not re-generate the configs.
-if [[ -n "$NUM_WORKERS" && -n "$GPU_IDS" && "$VARIANT_TYPE" == "gpu" ]]; then
-  # Count the number of GPU IDs provided
-  IFS=',' read -ra GPU_ID_ARRAY <<< "$GPU_IDS"
-  for i in "${GPU_ID_ARRAY[@]}"; do
+if [[ -n "$NUM_WORKERS" && "$VARIANT_TYPE" == "gpu" ]]; then
+  if [[ -n ${GPU_IDS:-} ]]; then
+      WORKER_IDS=($(echo "$GPU_IDS" | tr ',' ' '))
+  else
+      WORKER_IDS=($(seq 0 $((NUM_WORKERS - 1))))
+  fi
+  for i in "${WORKER_IDS[@]}"; do
     duplicate_worker_configs $i
   done
 fi

presto/scripts/run_benchmark.sh

@@ -30,6 +30,7 @@ OPTIONS:
                             stored inside a directory under the --output-dir path with a name matching the tag name.
                             Tags must contain only alphanumeric and underscore characters.
     -p, --profile           Enable profiling of benchmark queries.
+    --profile-script-path   Path to profiler functions script (default: ./profiler_functions.sh).
     --skip-drop-cache       Skip dropping system caches before each benchmark query (dropped by default).
     -m, --metrics           Collect detailed metrics from Presto REST API after each query.
                             Metrics are stored in query-specific directories.
@@ -147,6 +148,15 @@ parse_args() {
         PROFILE=true
         shift
         ;;
+      --profile-script-path)
+        if [[ -n $2 ]]; then
+          PROFILE_SCRIPT_PATH=$2
+          shift 2
+        else
+          echo "Error: --profile-script-path requires a value"
+          exit 1
+        fi
+        ;;
       --skip-drop-cache)
         SKIP_DROP_CACHE=true
         shift
@@ -218,7 +228,10 @@ if [[ -n ${TAG} ]]; then
 fi
 
 if [[ "${PROFILE}" == "true" ]]; then
-  PYTEST_ARGS+=("--profile --profile-script-path $(readlink -f ./profiler_functions.sh)")
+  if [[ -z "${PROFILE_SCRIPT_PATH:-}" ]]; then
+    PROFILE_SCRIPT_PATH="$(readlink -f ./profiler_functions.sh)"
+  fi
+  PYTEST_ARGS+=("--profile --profile-script-path ${PROFILE_SCRIPT_PATH}")
 fi
 
 if [[ "${METRICS}" == "true" ]]; then

presto/scripts/setup_benchmark_helper_check_instance_and_parse_args.sh

@@ -31,6 +31,7 @@ OPTIONS:
     -s, --schema-name                   Name of the schema that will contain the created tables.
     -d, --data-dir-name                 Name of the directory inside the PRESTO_DATA_DIR path for the benchmark data.
     --skip-analyze-tables               Skip analyzing tables after creating them. Default is to analyze tables.
+    --no-docker				Skip the setup/teardown steps that require docker.
     $SCRIPT_EXTRA_OPTIONS_DESCRIPTION
 
 EXAMPLES:
@@ -49,6 +50,7 @@ fi
 # Compute the directory where this script resides (if not already set by caller)
 SCRIPT_DIR="${SCRIPT_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)}"
 
+NO_DOCKER=false
 SKIP_ANALYZE_TABLES=false
 parse_args() {
   while [[ $# -gt 0 ]]; do
@@ -88,6 +90,10 @@ parse_args() {
         SKIP_ANALYZE_TABLES=true
         shift
         ;;
+      --no-docker)
+        NO_DOCKER=true
+        shift
+        ;;
       *)
         SCRIPT_EXTRA_OPTIONS_UNKNOWN_ARG=true
         if [[ -n $SCRIPT_EXTRA_OPTIONS_PARSER ]]; then

presto/scripts/setup_benchmark_tables.sh

@@ -32,11 +32,14 @@ function cleanup() {
 
 trap cleanup EXIT
 
-"${SCRIPT_DIR}/start_native_cpu_presto.sh"
-
-source "${SCRIPT_DIR}/common_functions.sh"
+# These scripts are used in some non-docker environments, so provide the option to skip
+# the docker setup/teardown.
+if [[ -z "$NO_DOCKER" ]]; then
+  "${SCRIPT_DIR}/start_native_cpu_presto.sh"
+  source "${SCRIPT_DIR}/common_functions.sh"
+  wait_for_worker_node_registration
+fi
 
-wait_for_worker_node_registration
 
 "${SCRIPT_DIR}/../../scripts/run_py_script.sh" -p $SCHEMA_GEN_SCRIPT_PATH \
                                --benchmark-type $BENCHMARK_TYPE \
@@ -53,4 +56,6 @@ if [[ "$SKIP_ANALYZE_TABLES" == "false" ]]; then
   "${SCRIPT_DIR}/analyze_tables.sh" -s $SCHEMA_NAME
 fi
 
-"${SCRIPT_DIR}/stop_presto.sh"
+if [[ -z "$NO_DOCKER" ]]; then
+  "${SCRIPT_DIR}/stop_presto.sh"
+fi

presto/scripts/start_presto_helper.sh

@@ -19,7 +19,12 @@ fi
 SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 
 # Get the root of the git repository
-REPO_ROOT="$(git -C "${SCRIPT_DIR}" rev-parse --show-toplevel)"
+if command -v git &> /dev/null; then
+  REPO_ROOT="$(git -C "${SCRIPT_DIR}" rev-parse --show-toplevel)"
+else
+  REPO_ROOT="$SCRIPT_DIR/../.."
+fi
+
 
 # Validate sibling repos
 if [[ "$VARIANT_TYPE" == "java" ]]; then
@@ -88,12 +93,6 @@ else
   echo "Internal error: unexpected VARIANT_TYPE value: $VARIANT_TYPE"
 fi
 
-# Default GPU_IDS if NUM_WORKERS is set but GPU_IDS is not
-if [[ -n $NUM_WORKERS && -z $GPU_IDS ]]; then
-  # Generate default GPU IDs: 0,1,2,...,N-1
-  export GPU_IDS=$(seq -s, 0 $((NUM_WORKERS - 1)))
-fi
-
 "${SCRIPT_DIR}/stop_presto.sh"
 
 "${SCRIPT_DIR}/generate_presto_config.sh"