diff --git a/presto/docker/docker-compose.native-cpu.yml b/presto/docker/docker-compose.native-cpu.yml index 4ecefe25..b249f16d 100644 --- a/presto/docker/docker-compose.native-cpu.yml +++ b/presto/docker/docker-compose.native-cpu.yml @@ -15,6 +15,7 @@ services: service: presto-base-native-worker container_name: presto-native-worker-cpu image: presto-native-worker-cpu:${PRESTO_IMAGE_TAG:?PRESTO_IMAGE_TAG must be set} + ulimits: {nofile: {soft: 500000, hard: 500000}} cap_add: - SYS_NICE build: diff --git a/presto/docker/docker-compose/template/docker-compose.native-cpu.yml.jinja b/presto/docker/docker-compose/template/docker-compose.native-cpu.yml.jinja new file mode 100644 index 00000000..d07d6dae --- /dev/null +++ b/presto/docker/docker-compose/template/docker-compose.native-cpu.yml.jinja @@ -0,0 +1,64 @@ +x-presto-native-worker-cpu: &cpu_worker_base + extends: + file: ../../docker-compose.common.yml + service: presto-base-native-worker + image: presto-native-worker-cpu:${PRESTO_IMAGE_TAG} + ulimits: {nofile: {soft: 500000, hard: 500000}} + # SYS_NICE is required so that numactl --cpunodebind / --membind work inside + # the container (needed for NUMA pinning of multi-worker CPU setups). + cap_add: + - SYS_NICE + build: + args: + - GPU=OFF + entrypoint: ["bash", "/opt/launch_presto_servers.sh"] + command: [] + depends_on: + - presto-coordinator + volumes: + # Mount /sys/devices/system/node so numactl inside the container can read + # the host NUMA topology when deciding memory/cpu bindings. + - /sys/devices/system/node:/sys/devices/system/node + +services: + presto-coordinator: + extends: + file: ../../docker-compose.common.yml + service: presto-base-coordinator + volumes: + - ../../config/generated/cpu/etc_common:/opt/presto-server/etc + - ../../config/generated/cpu/etc_coordinator/config_native.properties:/opt/presto-server/etc/config.properties + - ../../config/generated/cpu/etc_coordinator/node.properties:/opt/presto-server/etc/node.properties + - ../../config/generated/cpu/etc_coordinator/catalog/hive.properties:/opt/presto-server/etc/catalog/hive.properties + + {% if workers|length > 1 %} + # Multi-worker CPU: one container per worker, each pinned to a NUMA node. + # NUMA_NODE is read by launch_presto_servers.sh which wraps presto_server + # with `numactl --cpunodebind=N --membind=N`. + {% for worker in workers %} + presto-native-worker-cpu-{{ worker.id }}: + <<: *cpu_worker_base + container_name: presto-native-worker-cpu-{{ worker.id }} + environment: + WORKER_ID: {{ worker.id }} + NUMA_NODE: {{ worker.numa_node }} + SERVER_START_TIMESTAMP: ${SERVER_START_TIMESTAMP:-} + volumes: + - ../../config/generated/cpu/etc_common:/opt/presto-server/etc + - ../../config/generated/cpu/etc_worker_{{ worker.id }}/node.properties:/opt/presto-server/etc/node.properties + - ../../config/generated/cpu/etc_worker_{{ worker.id }}/config_native.properties:/opt/presto-server/etc/config.properties + - ../../config/generated/cpu/etc_worker_{{ worker.id }}/catalog/hive.properties:/opt/presto-server/etc/catalog/hive.properties + {% endfor %} + {% else %} + # Single-worker CPU: one container spanning all NUMA nodes via --interleave=all. + presto-native-worker-cpu: + <<: *cpu_worker_base + container_name: presto-native-worker-cpu + environment: + SERVER_START_TIMESTAMP: ${SERVER_START_TIMESTAMP:-} + volumes: + - ../../config/generated/cpu/etc_common:/opt/presto-server/etc + - ../../config/generated/cpu/etc_worker/node.properties:/opt/presto-server/etc/node.properties + - ../../config/generated/cpu/etc_worker/config_native.properties:/opt/presto-server/etc/config.properties + - ../../config/generated/cpu/etc_worker/catalog/hive.properties:/opt/presto-server/etc/catalog/hive.properties + {% endif %} diff --git a/presto/docker/launch_presto_servers.sh b/presto/docker/launch_presto_servers.sh index 2494b499..69e1cbb2 100644 --- a/presto/docker/launch_presto_servers.sh +++ b/presto/docker/launch_presto_servers.sh @@ -13,8 +13,12 @@ mkdir -p "${LOGS_DIR}" ETC_BASE="/opt/presto-server/etc" # Resolve the NUMA node for a worker and launch presto_server pinned to it. -# For GPU workers: pins to the NUMA node closest to the GPU via nvidia-smi topology. -# For CPU workers: interleaves memory across all NUMA nodes via numactl (requires SYS_NICE). +# Precedence: +# 1. NUMA_NODE env var: explicit bind (used by multi-worker CPU where each +# container represents one NUMA socket). +# 2. nvidia-smi topology: pins to the NUMA node closest to the worker's GPU. +# 3. numactl fallback: --interleave=all across all NUMA nodes (single- +# container CPU where one worker spans both sockets). # $1 — GPU ID (or 0 for CPU single-worker) # $2 — etc-dir path for this instance launch_worker() { @@ -24,7 +28,10 @@ launch_worker() { local launcher=() local cuda_env=() - if command -v nvidia-smi &> /dev/null; then + if [[ -n "${NUMA_NODE:-}" ]]; then + echo "NUMA_NODE=${NUMA_NODE} -- launching with numactl --cpunodebind=${NUMA_NODE} --membind=${NUMA_NODE}" + launcher=(numactl --cpunodebind="${NUMA_NODE}" --membind="${NUMA_NODE}") + elif command -v nvidia-smi &> /dev/null; then local topo topo=$(nvidia-smi topo -C -M -i "$worker_id") echo "$topo" @@ -60,11 +67,13 @@ launch_worker() { env "${cuda_env[@]}" "${launcher[@]}" presto_server --etc-dir="$etc_dir" >> "${log_file}" 2>&1 & } -# No args → single worker using CUDA_VISIBLE_DEVICES (default 0), shared config dir. +# No args → single worker. WORKER_ID env var (set by multi-worker CPU compose +# where each container is one logical worker) takes precedence; otherwise fall +# back to CUDA_VISIBLE_DEVICES for GPU runs, defaulting to 0. # With args → one worker per GPU ID, each with its own config dir (etc). if [ $# -eq 0 ]; then # Single worker mode. - launch_worker "${CUDA_VISIBLE_DEVICES:-0}" "${ETC_BASE}/" + launch_worker "${WORKER_ID:-${CUDA_VISIBLE_DEVICES:-0}}" "${ETC_BASE}/" else # Multi-worker single-container mode. Each GPU ID is an argument. for gpu_id in "$@"; do diff --git a/presto/docker/presto_profiling_wrapper.sh b/presto/docker/presto_profiling_wrapper.sh index 637c1360..a987fc2a 100644 --- a/presto/docker/presto_profiling_wrapper.sh +++ b/presto/docker/presto_profiling_wrapper.sh @@ -5,7 +5,7 @@ set -e if [[ "$PROFILE" == "ON" ]]; then - mkdir /presto_profiles + mkdir -p /presto_profiles if [[ -z $PROFILE_ARGS ]]; then PROFILE_ARGS="-t nvtx,cuda" diff --git a/presto/scripts/generate_presto_config.sh b/presto/scripts/generate_presto_config.sh index 611b4393..b062c8ce 100755 --- a/presto/scripts/generate_presto_config.sh +++ b/presto/scripts/generate_presto_config.sh @@ -34,7 +34,7 @@ fi # 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" + echo "Duplicating worker configs for worker ID $worker_id" local worker_config="${CONFIG_DIR}/etc_worker_${worker_id}" local coord_config="${CONFIG_DIR}/etc_coordinator" local worker_native_config="${worker_config}/config_native.properties" @@ -69,6 +69,18 @@ NPROC=$(nproc) # lsmem will report in SI. Make sure we get values in GB. RAM_GB=$(lsmem -b | grep "Total online memory" | awk '{print int($4 / (1024*1024*1024)); }') +# Detect host NUMA / SMT topology. Used below to auto-tune CPU workers so the +# same scripts produce sensible configs on any host (not just 0374). +# `lscpu -p` emits one machine-parseable row per logical CPU; each row has +# fields `CPU,Core,Socket,Node,...`. Counting unique (Socket,Core) pairs gives +# the physical-core count regardless of SMT state or lscpu locale strings. +NUMA_NODES=$(ls -d /sys/devices/system/node/node* 2>/dev/null | wc -l) +[[ ${NUMA_NODES} -lt 1 ]] && NUMA_NODES=1 +PHYSICAL_CORES=$(lscpu -p 2>/dev/null | grep -v '^#' | awk -F, '{print $3","$2}' | sort -u | wc -l) +[[ -z "${PHYSICAL_CORES}" || "${PHYSICAL_CORES}" -lt 1 ]] && PHYSICAL_CORES=${NPROC} +SMT_RATIO=$(( NPROC / PHYSICAL_CORES )) +[[ ${SMT_RATIO} -lt 1 ]] && SMT_RATIO=1 + # variant-specific behavior # for GPU you must set vcpu_per_worker to a small number, not the CPU count if [[ -z ${VARIANT_TYPE} || ! ${VARIANT_TYPE} =~ ^(cpu|gpu|java)$ ]]; then @@ -162,8 +174,17 @@ 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" && "$VARIANT_TYPE" == "gpu" ]]; then - if [[ -n ${GPU_IDS:-} ]]; then +# +# GPU always gets per-worker config dirs (even -w 1) because its compose +# template mounts etc_worker_ in both multi-worker and combined-container +# layouts. GPU IDs are explicit (via GPU_IDS) since each worker is bound to +# a specific device. +# +# CPU only needs per-worker dirs when multi-worker (-w N > 1); the -w 1 CPU +# path mounts etc_worker directly, matching historical single-container +# behaviour. CPU workers use 0-indexed sequential IDs. +if [[ -n "$NUM_WORKERS" && ( "$VARIANT_TYPE" == "gpu" || ( "$VARIANT_TYPE" == "cpu" && "$NUM_WORKERS" -gt 1 ) ) ]]; then + if [[ "$VARIANT_TYPE" == "gpu" && -n ${GPU_IDS:-} ]]; then WORKER_IDS=($(echo "$GPU_IDS" | tr ',' ' ')) else WORKER_IDS=($(seq 0 $((NUM_WORKERS - 1)))) @@ -172,3 +193,122 @@ if [[ -n "$NUM_WORKERS" && "$VARIANT_TYPE" == "gpu" ]]; then duplicate_worker_configs $i done fi + +# Reconcile single-node-execution-enabled and cudf.exchange on every start. +# +# duplicate_worker_configs flips both to their multi-worker values when +# NUM_WORKERS > 1 (single-node-execution-enabled=false on coord + workers, +# cudf.exchange=true on workers), and never reverts them for NUM_WORKERS == 1. +# Switching from -w N>1 to -w 1 without `--overwrite-config` therefore leaves +# stale multi-worker values in place. Concretely, stale +# single-node-execution-enabled=false on the coord makes the planner keep +# generating multi-stage distributed plans, which a lone worker then executes +# with HTTP-exchange roundtrips between stages — measured 5x regression on +# TPC-H Q17 for both -w 1 CPU and -w 1 GPU after a multi-worker run. +# +# Always set both to match the current NUM_WORKERS so toggling worker count +# doesn't require --overwrite-config. +if [[ -n "$NUM_WORKERS" && ( "$VARIANT_TYPE" == "gpu" || "$VARIANT_TYPE" == "cpu" ) ]]; then + if [[ "$NUM_WORKERS" -gt 1 ]]; then + SINGLE_NODE_EXECUTION="false" + CUDF_EXCHANGE="true" + else + SINGLE_NODE_EXECUTION="true" + CUDF_EXCHANGE="false" + fi + for cfg in \ + "${CONFIG_DIR}/etc_coordinator/config_native.properties" \ + "${CONFIG_DIR}"/etc_worker*/config_native.properties; do + [[ -f "$cfg" ]] || continue + sed -i "s/^single-node-execution-enabled=.*/single-node-execution-enabled=${SINGLE_NODE_EXECUTION}/" "$cfg" + done + for cfg in "${CONFIG_DIR}"/etc_worker*/config_native.properties; do + [[ -f "$cfg" ]] || continue + sed -i "s/^cudf.exchange=.*/cudf.exchange=${CUDF_EXCHANGE}/" "$cfg" + done +fi + +# CPU auto-tuning keyed on NUM_WORKERS + detected host topology. +# Runs idempotently on every start so manual sed-per-run isn't required, and +# switching between `-w 1` and `-w 2` just needs `stop_presto.sh && start_native_cpu_presto.sh -w N`. +# Any value can be overridden via env var. +# +# Logic: +# -w 1 -> container spans all NUMA nodes (launcher uses --interleave=all). +# Give it the full host memory; drivers = min(NPROC, 255) to max SMT +# up to Velox's uint8_t cap on parallel hash-build tables. +# -w N -> each container pins to one NUMA node via NUMA_NODE env var. +# Memory budget = per-node RAM; drivers = min(threads-per-node, 255) +# so each worker has 1 driver per hw thread on its socket, no cross- +# container oversubscription of the NUMA's cores. +# Cache headroom (system-memory-gb - query-memory-gb) defaults to ~30% of the +# worker's memory envelope, enough for the async data cache to hold a +# meaningful chunk of SF1000 TPC-H. +if [[ "${VARIANT_TYPE}" == "cpu" ]]; then + : "${CPU_ASYNC_DATA_CACHE:=true}" + if [[ "${NUM_WORKERS:-1}" -gt 1 ]]; then + THREADS_PER_WORKER=$(( NPROC / NUM_WORKERS )) + RAM_PER_WORKER_GB=$(( RAM_GB / NUM_WORKERS )) + : "${CPU_SYSTEM_MEM_GB:=$(( RAM_PER_WORKER_GB - 35 ))}" + : "${CPU_DRIVERS:=$(( THREADS_PER_WORKER < 255 ? THREADS_PER_WORKER : 255 ))}" + # Multi-worker bottlenecks on HTTP exchange backpressure for shuffle-heavy + # queries (Q9/Q18 time out at 30m with the 32MB defaults). 16x larger + # buffers unlock those queries on SF1000. + : "${CPU_EXCHANGE_BUFFER:=512MB}" + : "${CPU_SINK_BUFFER:=512MB}" + else + : "${CPU_SYSTEM_MEM_GB:=${RAM_GB}}" + : "${CPU_DRIVERS:=$(( NPROC < 255 ? NPROC : 255 ))}" + # Single-worker has no inter-worker shuffle, so the exchange buffers + # only matter for in-process exchange. Use Velox's defaults (32MB) — + # bigger buffers add latency on small shuffles for no benefit here. + : "${CPU_EXCHANGE_BUFFER:=32MB}" + : "${CPU_SINK_BUFFER:=32MB}" + fi + # ~30% of the worker's memory envelope reserved for the Velox async data cache + : "${CPU_QUERY_MEM_GB:=$(( CPU_SYSTEM_MEM_GB * 70 / 100 ))}" + CPU_SYSTEM_MEM_LIMIT_GB=$(( CPU_SYSTEM_MEM_GB + 30 )) + + echo "CPU auto-tune: NUM_WORKERS=${NUM_WORKERS:-1} NPROC=${NPROC} NUMA_NODES=${NUMA_NODES} PHYSICAL_CORES=${PHYSICAL_CORES} SMT_RATIO=${SMT_RATIO}" + echo " system-memory-gb=${CPU_SYSTEM_MEM_GB} query-memory-gb=${CPU_QUERY_MEM_GB} task.max-drivers-per-task=${CPU_DRIVERS} async-data-cache-enabled=${CPU_ASYNC_DATA_CACHE}" + echo " exchange.max-buffer-size=${CPU_EXCHANGE_BUFFER} sink.max-buffer-size=${CPU_SINK_BUFFER}" + + # set_or_append : idempotent — replaces existing + # `key=...` line if present, otherwise appends. Used here because the + # buffer-size keys aren't in the pbench template (Velox falls back to + # built-in defaults), so on first apply we have to append; on subsequent + # applies we replace. + set_or_append() { + local key=$1 value=$2 cfg=$3 + if grep -q "^${key}=" "$cfg"; then + sed -i "s|^${key}=.*|${key}=${value}|" "$cfg" + else + echo "${key}=${value}" >> "$cfg" + fi + } + + # Apply to every CPU worker config dir: etc_worker (single-worker + template + # for duplicate_worker_configs) plus etc_worker_ (multi-worker instances). + for worker_dir in "${CONFIG_DIR}"/etc_worker*/; do + cfg="${worker_dir}config_native.properties" + [[ -f "$cfg" ]] || continue + sed -i "s/^system-memory-gb=.*/system-memory-gb=${CPU_SYSTEM_MEM_GB}/" "$cfg" + sed -i "s/^system-mem-limit-gb=.*/system-mem-limit-gb=${CPU_SYSTEM_MEM_LIMIT_GB}/" "$cfg" + sed -i "s/^query-memory-gb=.*/query-memory-gb=${CPU_QUERY_MEM_GB}/" "$cfg" + sed -i "s|^query.max-memory-per-node=.*|query.max-memory-per-node=${CPU_QUERY_MEM_GB}GB|" "$cfg" + sed -i "s/^task.max-drivers-per-task=.*/task.max-drivers-per-task=${CPU_DRIVERS}/" "$cfg" + sed -i "s/^async-data-cache-enabled=.*/async-data-cache-enabled=${CPU_ASYNC_DATA_CACHE}/" "$cfg" + set_or_append "exchange.max-buffer-size" "${CPU_EXCHANGE_BUFFER}" "$cfg" + set_or_append "sink.max-buffer-size" "${CPU_SINK_BUFFER}" "$cfg" + done + + # hive.file-splittable flip for CPU. Put outside the regen gate so restarts + # without a full regen still land the right value; idempotent because the + # pattern only matches the untouched GPU-oriented template value. + for hive_cfg in \ + "${CONFIG_DIR}/etc_coordinator/catalog/hive.properties" \ + "${CONFIG_DIR}"/etc_worker*/catalog/hive.properties; do + [[ -f "$hive_cfg" ]] || continue + sed -i 's/^hive.file-splittable=false/hive.file-splittable=true/' "$hive_cfg" + done +fi diff --git a/presto/scripts/start_presto_helper.sh b/presto/scripts/start_presto_helper.sh index 39480c31..60d64676 100755 --- a/presto/scripts/start_presto_helper.sh +++ b/presto/scripts/start_presto_helper.sh @@ -126,6 +126,15 @@ elif [[ "$VARIANT_TYPE" == "cpu" ]]; then else DOCKER_COMPOSE_FILE="native-cpu" fi + # When the multi-worker CPU jinja template is in use (NUM_WORKERS>1, + # not single-container), the rendered file names per-worker services + # presto-native-worker-cpu-{0..N-1}. All workers share the same image, + # so picking worker 0 as the build target is sufficient. sccache uses a + # static file with only one CPU service, so skip the rename in that case. + CPU_TEMPLATE_PATH="${SCRIPT_DIR}/../docker/docker-compose/template/docker-compose.$DOCKER_COMPOSE_FILE.yml.jinja" + if [[ -f "$CPU_TEMPLATE_PATH" && -n "$NUM_WORKERS" && "$NUM_WORKERS" -gt 1 && "$SINGLE_CONTAINER" == "false" ]]; then + CPU_WORKER_SERVICE="presto-native-worker-cpu-0" + fi conditionally_add_build_target $CPU_WORKER_IMAGE $CPU_WORKER_SERVICE "worker|w" elif [[ "$VARIANT_TYPE" == "gpu" ]]; then DOCKER_COMPOSE_FILE="native-gpu" @@ -181,9 +190,12 @@ elif [[ "$ALL_CUDA_ARCHS" == "true" ]]; then fi DOCKER_COMPOSE_FILE_PATH="${SCRIPT_DIR}/../docker/docker-compose.$DOCKER_COMPOSE_FILE.yml" -# For GPU, the docker-compose file is a Jinja template. Render it before any docker compose operations. -if [[ "$VARIANT_TYPE" == "gpu" ]]; then - TEMPLATE_PATH="${SCRIPT_DIR}/../docker/docker-compose/template/docker-compose.$DOCKER_COMPOSE_FILE.yml.jinja" +# GPU always renders from a Jinja template (per-GPU config mounts are parameterised). +# CPU renders from a Jinja template only when one exists for the current compose +# file name — the default "native-cpu" has a template to support per-NUMA +# multi-worker layouts, but "native-cpu.sccache" keeps using its static file. +TEMPLATE_PATH="${SCRIPT_DIR}/../docker/docker-compose/template/docker-compose.$DOCKER_COMPOSE_FILE.yml.jinja" +if [[ "$VARIANT_TYPE" == "gpu" || ( "$VARIANT_TYPE" == "cpu" && -f "$TEMPLATE_PATH" ) ]]; then RENDERED_DIR="${SCRIPT_DIR}/../docker/docker-compose/generated" mkdir -p "$RENDERED_DIR" RENDERED_PATH="$RENDERED_DIR/docker-compose.$DOCKER_COMPOSE_FILE.rendered.yml" @@ -191,8 +203,8 @@ if [[ "$VARIANT_TYPE" == "gpu" ]]; then LOCAL_NUM_WORKERS="${NUM_WORKERS:-0}" RENDER_SCRIPT_PATH=$(readlink -f "${SCRIPT_DIR}/../../template_rendering/render_docker_compose_template.py") - RENDER_ARGS="--template-path $TEMPLATE_PATH --output-path $RENDERED_PATH --num-workers $NUM_WORKERS --single-container $SINGLE_CONTAINER --kvikio-threads $KVIKIO_THREADS --sccache $ENABLE_SCCACHE" - if [[ -n $GPU_IDS ]]; then + RENDER_ARGS="--template-path $TEMPLATE_PATH --output-path $RENDERED_PATH --num-workers $NUM_WORKERS --single-container $SINGLE_CONTAINER --kvikio-threads $KVIKIO_THREADS --sccache $ENABLE_SCCACHE --variant $VARIANT_TYPE" + if [[ "$VARIANT_TYPE" == "gpu" && -n $GPU_IDS ]]; then RENDER_ARGS="$RENDER_ARGS --gpu-ids $GPU_IDS" fi "${SCRIPT_DIR}/../../scripts/run_py_script.sh" -p "$RENDER_SCRIPT_PATH" $RENDER_ARGS @@ -237,3 +249,32 @@ fi # Start all services defined in the rendered docker-compose file. docker compose -f $DOCKER_COMPOSE_FILE_PATH up -d + +# Optional: block until ${NUM_WORKERS} workers have registered with the +# coordinator. Useful when a subsequent step (benchmark run, post-start +# setup) needs the cluster to be fully up. Times out after +# WAIT_FOR_WORKERS_TIMEOUT seconds (default 120) so a stuck worker +# doesn't hang the shell indefinitely. +if [[ "$WAIT_FOR_WORKERS" == "true" ]]; then + echo "Waiting for ${NUM_WORKERS} worker(s) to register with coordinator (timeout: ${WAIT_FOR_WORKERS_TIMEOUT}s)..." + deadline=$(( $(date +%s) + WAIT_FOR_WORKERS_TIMEOUT )) + workers=0 + while [[ $(date +%s) -lt $deadline ]]; do + workers=$(python3 - <<'PY' +import json, urllib.request +try: + with urllib.request.urlopen("http://localhost:8080/v1/node", timeout=1) as r: + print(len(json.load(r))) +except Exception: + print(0) +PY + ) + echo " ${workers}/${NUM_WORKERS} worker(s) registered" + [[ "$workers" -ge "$NUM_WORKERS" ]] && break + sleep 5 + done + if [[ "$workers" -lt "$NUM_WORKERS" ]]; then + echo "ERROR: only ${workers}/${NUM_WORKERS} workers registered after ${WAIT_FOR_WORKERS_TIMEOUT}s" >&2 + exit 1 + fi +fi diff --git a/presto/scripts/start_presto_helper_parse_args.sh b/presto/scripts/start_presto_helper_parse_args.sh index 6b868b94..3275649e 100644 --- a/presto/scripts/start_presto_helper_parse_args.sh +++ b/presto/scripts/start_presto_helper_parse_args.sh @@ -76,6 +76,8 @@ ENABLE_SCCACHE=false SCCACHE_AUTH_DIR="${SCCACHE_AUTH_DIR:-$HOME/.sccache-auth}" SCCACHE_ENABLE_DIST=false SCCACHE_VERSION="${SCCACHE_VERSION:-latest}" +WAIT_FOR_WORKERS=false +WAIT_FOR_WORKERS_TIMEOUT=120 parse_args() { while [[ $# -gt 0 ]]; do case $1 in @@ -206,6 +208,19 @@ parse_args() { SCCACHE_ENABLE_DIST=true shift ;; + --wait) + WAIT_FOR_WORKERS=true + shift + ;; + --wait-timeout) + if [[ -n $2 ]]; then + WAIT_FOR_WORKERS_TIMEOUT=$2 + shift 2 + else + echo "Error: --wait-timeout requires a value (seconds)" + exit 1 + fi + ;; *) echo "Error: Unknown argument $1" print_help diff --git a/presto/scripts/stop_presto.sh b/presto/scripts/stop_presto.sh index 37628883..c1fe0dc3 100755 --- a/presto/scripts/stop_presto.sh +++ b/presto/scripts/stop_presto.sh @@ -4,6 +4,11 @@ set -e +# Ignore empty var warnings from docker compose down +: "${PROFILE:=}" +: "${PROFILE_ARGS:=}" +export PROFILE PROFILE_ARGS + # Compute the directory where this script resides SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" @@ -12,12 +17,35 @@ if [ -z "${PRESTO_IMAGE_TAG}" ]; then fi GPU_FILE="${SCRIPT_DIR}/../docker/docker-compose/generated/docker-compose.native-gpu.rendered.yml" +CPU_RENDERED_FILE="${SCRIPT_DIR}/../docker/docker-compose/generated/docker-compose.native-cpu.rendered.yml" JAVA_FILE="${SCRIPT_DIR}/../docker/docker-compose.java.yml" CPU_FILE="${SCRIPT_DIR}/../docker/docker-compose.native-cpu.yml" -# Bring down each variant independently to avoid path resolution issues when combining files -docker compose -f "$JAVA_FILE" down -docker compose -f "$CPU_FILE" down -if [ -f "$GPU_FILE" ]; then - docker compose -f "$GPU_FILE" down +# Tear down whichever variant is currently running (if any). The coordinator +# is shared via `extends` across every variant's compose file, so calling +# `down` on a non-active variant would remove it and orphan the real workers +# on the shared network — hence the single-file policy below. Uses `docker +# ps -a` so we also catch containers in Created or Exited state left over +# from a half-finished prior run. +case "$(docker ps -a --format '{{.Names}}' 2>/dev/null | grep -m1 '^presto-native-worker\|^presto-java-worker')" in + presto-native-worker-gpu*) ACTIVE="$GPU_FILE" ;; + presto-native-worker-cpu-[0-9]*) ACTIVE="$CPU_RENDERED_FILE" ;; + presto-native-worker-cpu) [ -f "$CPU_RENDERED_FILE" ] && ACTIVE="$CPU_RENDERED_FILE" || ACTIVE="$CPU_FILE" ;; + presto-java-worker*) ACTIVE="$JAVA_FILE" ;; + *) ACTIVE="" ;; +esac + +if [ -n "$ACTIVE" ] && [ -f "$ACTIVE" ]; then + docker compose -f "$ACTIVE" down +fi + +# Safety net: force-remove any lingering presto-* containers. `docker +# compose down` only removes containers whose `com.docker.compose.project` +# label matches the compose file's project name, so an orphan coordinator +# started by a different project (e.g. the rendered-file project when +# we're now acting on the static file) is silently skipped. `docker rm -f` +# has no such constraint. +orphans=$(docker ps -a --format '{{.Names}}' 2>/dev/null | grep '^presto-' || true) +if [ -n "$orphans" ]; then + echo "$orphans" | xargs -r docker rm -f fi diff --git a/template_rendering/render_docker_compose_template.py b/template_rendering/render_docker_compose_template.py index 080db9d6..8e3f92f8 100644 --- a/template_rendering/render_docker_compose_template.py +++ b/template_rendering/render_docker_compose_template.py @@ -4,9 +4,27 @@ import argparse import os +import re import sys +def detect_numa_nodes(): + """Return list of NUMA node IDs visible on this host. + + Reads /sys/devices/system/node/node entries. Falls back to [0] when + the sysfs layout is unavailable (non-Linux, minimal container, etc.). + """ + node_dir = "/sys/devices/system/node" + if not os.path.isdir(node_dir): + return [0] + nodes = [] + for entry in sorted(os.listdir(node_dir)): + m = re.match(r"^node(\d+)$", entry) + if m: + nodes.append(int(m.group(1))) + return nodes or [0] + + def parse_args(): parser = argparse.ArgumentParser(description="Render a Docker Compose template") @@ -52,6 +70,16 @@ def str_to_bool(value: str) -> bool: required=False, help="Enable sccache build secrets in the rendered compose file.", ) + parser.add_argument( + "--variant", + type=str, + default="gpu", + choices=["gpu", "cpu"], + dest="variant", + required=False, + help="Which variant this template describes. 'cpu' uses NUMA-node assignment; " + "'gpu' uses per-GPU assignment via --gpu-ids.", + ) return parser.parse_args() @@ -82,11 +110,24 @@ def main() -> int: ) template = env.get_template(os.path.basename(parsed_args.template_path)) - # If GPU IDs are provided, use them; otherwise default to range - if gpu_ids: - workers = gpu_ids + # Build the worker list. + # - GPU variant (default): plain list of GPU IDs. Preserves the existing + # contract the GPU template expects (worker loop variable is the GPU id). + # - CPU variant: list of dicts with {id, numa_node}. NUMA assignment is + # round-robin across the NUMA nodes detected on the host so that with + # --num-workers equal to the node count each worker lands on its own + # socket. With fewer workers than nodes, leading nodes are used. + if parsed_args.variant == "gpu": + if gpu_ids: + workers = gpu_ids + else: + workers = list(range(max(0, parsed_args.num_workers))) else: - workers = list(range(max(0, parsed_args.num_workers))) + numa_nodes = detect_numa_nodes() + workers = [ + {"id": i, "numa_node": numa_nodes[i % len(numa_nodes)]} + for i in range(max(0, parsed_args.num_workers)) + ] rendered = template.render( num_workers=parsed_args.num_workers, @@ -94,6 +135,7 @@ def main() -> int: single_container=parsed_args.single_container, kvikio_threads=parsed_args.kvikio_threads, sccache=parsed_args.sccache, + variant=parsed_args.variant, ) os.makedirs(os.path.dirname(parsed_args.output_path), exist_ok=True)