Misiug/multi node docker containers

.gitignore

@@ -18,3 +18,4 @@ benchmark_output
 
 # Generated Presto Config
 presto/docker/config/generated/
+presto/docker/docker-compose.native-gpu.yml

presto/scripts/generate_presto_config.sh

@@ -97,9 +97,9 @@ EOF
     # optimizer.default-filter-factor-enabled=true
     COORD_CONFIG="${CONFIG_DIR}/etc_coordinator/config_native.properties"
     sed -i 's/\#optimizer/optimizer/g' ${COORD_CONFIG}
-    
+
     # Adds a cluster tag for gpu variant
-    WORKER_CONFIG="${CONFIG_DIR}/etc_coordinator/config_native.properties"
+    WORKER_CONFIG="${CONFIG_DIR}/etc_worker/config_native.properties"
     echo "cluster-tag=native-gpu" >> ${WORKER_CONFIG}
   fi
 

presto/scripts/start_presto_helper.sh

@@ -66,7 +66,7 @@ function conditionally_add_build_target() {
     echo "Added $2 to the list of services to build because the '$BUILD_TARGET' build target was specified"
     BUILD_TARGET_ARG+=($2)
   fi
-} 
+}
 
 conditionally_add_build_target $COORDINATOR_IMAGE $COORDINATOR_SERVICE "coordinator|c"
 
@@ -130,4 +130,86 @@ if (( ${#BUILD_TARGET_ARG[@]} )); then
   ${BUILD_TARGET_ARG[@]}
 fi
 
+function generate_worker_config() {
+    local worker_id=$1
+    local worker_config="../docker/config/generated/gpu/etc_worker_${worker_id}"
+    rm -rf ${worker_config}
+    # We are duplicating the config generated by pbench,
+    # but need to modify it for multi-worker execution.
+    cp -r ../docker/config/generated/gpu/etc_worker ${worker_config}
+
+    # Disable single-execution mode.
+    sed -i "s+single-node-execution-enabled.*+single-node-execution-enabled=false+g" \
+        ${worker_config}/config_native.properties
+
+    # Give each worker a unique id.
+    sed -i "s+node\.id.*+node\.id=worker_${worker_id}+g" ${worker_config}/node.properties
+}
+
+function generate_worker_service() {
+    # These variables were created so that we have more fine-grain control over the naming
+    # of services and files (so we don't have to change legacy names for 1-worker setup)
+    local worker_suffix="$1"
+    local file_suffix="$2"
+    local worker_gpu="$3"
+
+    cat >> "$DOCKER_COMPOSE_FILE_PATH" <<YAML
+
+  presto-native-worker-gpu${worker_suffix}:
+    extends:
+      file: docker-compose.common.yml
+      service: presto-base-native-worker
+    container_name: presto-native-worker-gpu${worker_suffix}
+    image: presto-native-worker-gpu:latest
+    build:
+      args:
+        - GPU=ON
+    runtime: nvidia
+    environment:
+      - NVIDIA_VISIBLE_DEVICES=${worker_gpu}
+      - PROFILE=\${PROFILE}
+      - PROFILE_ARGS=\${PROFILE_ARGS}
+    depends_on:
+      - presto-coordinator
+    volumes:
+      - ./config/generated/gpu/etc_common:/opt/presto-server/etc
+      - ./config/generated/gpu/etc_worker${file_suffix}/node.properties:/opt/presto-server/etc/node.properties
+      - ./config/generated/gpu/etc_worker${file_suffix}/config_native.properties:/opt/presto-server/etc/config.properties
+YAML
+}
+
+function generate_coordinator_service() {
+    # Start the override file
+    cat > "$DOCKER_COMPOSE_FILE_PATH" <<YAML
+services:
+  presto-coordinator:
+    extends:
+      file: docker-compose.common.yml
+      service: presto-base-coordinator
+    volumes:
+      - ./config/generated/gpu/etc_common:/opt/presto-server/etc
+      - ./config/generated/gpu/etc_coordinator/config_native.properties:/opt/presto-server/etc/config.properties
+      - ./config/generated/gpu/etc_coordinator/node.properties:/opt/presto-server/etc/node.properties
+YAML
+}
+
+function generate_worker_compose() {
+    generate_coordinator_service
+
+    if [[ "$NUM_WORKERS" -gt "1" ]]; then
+        # Generate NUM_WORKERS services, each with one gpu.
+        for i in $(seq 0 $((NUM_WORKERS-1))); do
+            generate_worker_service "-$i" "_$i" "$i"
+            generate_worker_config $i
+        done
+    else
+        # Generate one worker service with access to all gpus
+        generate_worker_service "" "" "all"
+    fi
+}
+
+if [[ "$VARIANT_TYPE" == "gpu" ]]; then
+    generate_worker_compose
+fi
+
 docker compose -f $DOCKER_COMPOSE_FILE_PATH up -d