WIP: Common config and Dockerfile for NGC

Dockerfile

@@ -0,0 +1,48 @@
+# syntax = docker/dockerfile:1.0-experimental
+ARG BASE_IMAGE=nvidia/cuda:11.0-runtime-ubuntu20.04
+FROM ${BASE_IMAGE}
+
+ARG GROUP_ID=10000
+ARG USER_ID=10000
+ARG USER_GROUP_NAME=rapids
+ARG CONDA_ENV_FILE=conda/rapids-tpcx-bb.yml
+
+ENV PATH /conda/bin:$PATH
+ADD https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh miniconda.sh
+RUN bash miniconda.sh -f -b -p /conda && \
+    conda init bash && \
+    apt-get update && apt-get install -y git vim
+
+ENV CUDA_VERSION=${CUDA_VERSION}
+ENV NVIDIA_VISIBLE_DEVICES=${NVIDIA_VISIBLE_DEVICES}
+ENV hostname=${hostname}_cuda${CUDA_VERSION}-${DISTRO}
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV OMPI_MCA_opal_cuda_support=true
+
+#Workaround for http://nvbugs/2903712
+ENV IBV_DRIVERS=/usr/lib/libibverbs/libmlx5
+
+RUN groupadd --gid ${GROUP_ID} ${USER_GROUP_NAME} && \
+    useradd -g ${GROUP_ID} -u ${USER_ID} -ms /bin/bash ${USER_GROUP_NAME} && \
+    cat /conda/etc/profile.d/conda.sh >> /etc/profile && \
+    echo "conda activate tpcxbb" >> /etc/profile && \
+    cat /conda/etc/profile.d/conda.sh >> /home/${USER_GROUP_NAME}/.bashrc && \
+    echo "conda activate tpcxbb" >> /home/${USER_GROUP_NAME}/.bashrc && \
+    chgrp -R ${USER_GROUP_NAME} /conda && \
+    chmod -R g+rwx /conda
+
+COPY ${CONDA_ENV_FILE} /home/${USER_GROUP_NAME}/environment.yml
+RUN bin/bash -c "source /conda/etc/profile.d/conda.sh && \
+                 conda env create -f /home/${USER_GROUP_NAME}/environment.yml -n tpcxbb" && \
+    chmod -R g+w /conda && \
+    chown -R ${USER_GROUP_NAME}:${USER_GROUP_NAME} /conda/envs/tpcxbb
+
+COPY tpcx_bb /home/${USER_GROUP_NAME}/tpcx_bb
+RUN chown -R ${USER_GROUP_NAME}:${USER_GROUP_NAME} /home/${USER_GROUP_NAME}/tpcx_bb && \
+    /bin/bash -c "source /conda/etc/profile.d/conda.sh && \
+                  conda activate tpcxbb && \
+                  python -m pip install /home/${USER_GROUP_NAME}/tpcx_bb/."
+
+USER ${USER_GROUP_NAME}
+WORKDIR /home/${USER_GROUP_NAME}/

conda/rapids-tpcx-bb.yml

@@ -1,5 +1,5 @@
 channels:
-  - blazingsql-nightly/label/cuda10.2
+  - blazingsql-nightly
   - rapidsai-nightly
   - nvidia
   - pytorch-nightly
@@ -9,7 +9,7 @@ channels:
 
 dependencies:
   - python=3.7
-  - cudatoolkit=10.2
+  - cudatoolkit
   - cudf
   - rmm
   - dask-cuda

tpcx_bb/benchmark_runner.py

@@ -11,10 +11,11 @@ def get_qnum_from_filename(name):
     m = re.search("[0-9]{2}", name).group()
     return m
 
+base_path = os.path.split(os.path.abspath(__file__))[0]
 
 dask_qnums = [str(i).zfill(2) for i in range(1, 31)]
 # Not all queries are implemented with BSQL
-bsql_query_files = sorted(glob.glob("./queries/q*/t*_sql.py"))
+bsql_query_files = sorted(glob.glob(os.path.join( base_path, "queries/q*/t*_sql.py")))
 bsql_qnums = [get_qnum_from_filename(x.split("/")[-1]) for x in bsql_query_files]
 
 def load_query(qnum, fn):
@@ -24,7 +25,7 @@ def load_query(qnum, fn):
     loader.exec_module(mod)
     return mod.main
 
-if __name__ == "__main__":
+def main():
     from xbb_tools.cluster_startup import attach_to_cluster, import_query_libs
     from xbb_tools.utils import run_query, tpcxbb_argparser
 
@@ -41,14 +42,12 @@ def load_query(qnum, fn):
 
 
     config = tpcxbb_argparser()
-    include_blazing = config.get("benchmark_runner_include_bsql")
+    include_blazing = config.get("with_blazing")
     client, bc = attach_to_cluster(config, create_blazing_context=include_blazing)
 
     # Preload required libraries for queries on all workers
     client.run(import_query_libs)
 
-    base_path = os.getcwd()
-
     # Run Pure Dask Queries
     if len(dask_qnums) > 0:
         print("Pure Dask Queries")
@@ -93,3 +92,6 @@ def load_query(qnum, fn):
                 client.run_on_scheduler(gc.collect)
                 gc.collect()
                 time.sleep(3)
+
+if __name__ == "__main__":
+    main()

tpcx_bb/cluster_configuration/bsql-cluster-startup.sh

@@ -1,25 +1,29 @@
+#!/usr/bin/env bash
+
+set -x
+
 #IB, NVLINK, or TCP
-CLUSTER_MODE=$1
+CLUSTER_MODE=${CLUSTER_MODE:="IB"}
 USERNAME=$(whoami)
 
 MAX_SYSTEM_MEMORY=$(free -m | awk '/^Mem:/{print $2}')M
-DEVICE_MEMORY_LIMIT="25GB"
-POOL_SIZE="30GB"
+DEVICE_MEMORY_LIMIT=${DEVICE_MEMORY_LIMIT:="25GB"}
+POOL_SIZE=${POOL_SIZE:="30GB"}
 
 # Fill in your environment name and conda path on each node
-TPCX_BB_HOME="/home/$USERNAME/shared/tpcx-bb"
-CONDA_ENV_NAME="rapids-tpcx-bb"
-CONDA_ENV_PATH="/home/$USERNAME/conda/etc/profile.d/conda.sh"
+TPCX_BB_HOME=${TPCX_BB_HOME:="/home/$USERNAME/shared/tpcx-bb"}
+CONDA_ENV_NAME=${CONDA_ENV_NAME:="rapids-tpcx-bb"}
+CONDA_ENV_PATH=${CONDA_ENV_PATH:="/home/$USERNAME/conda/etc/profile.d/conda.sh"}
 
 # TODO: Unify interface/IP setting/getting for cluster startup
 # and scheduler file
-INTERFACE="ib0"
+INTERFACE=${INTERFACE:="ib0"}
 
 # TODO: Remove hard-coding of scheduler
-SCHEDULER=$(hostname)
-SCHEDULER_FILE=$TPCX_BB_HOME/tpcx_bb/cluster_configuration/example-cluster-scheduler.json
-LOGDIR="/tmp/tpcx-bb-dask-logs/"
-WORKER_DIR="/tmp/tpcx-bb-dask-workers/"
+SCHEDULER=${SCHEDULER:=$(hostname)}
+SCHEDULER_FILE=${SCHEDULER_FILE:="${TPCX_BB_HOME}/tpcx_bb/cluster_configuration/example-cluster-scheduler.json"}
+LOGDIR=${LOGDIR:="/tmp/tpcx-bb-dask-logs/"}
+WORKER_DIR=${WORKER_DIR:="/tmp/tpcx-bb-dask-workers/"}
 
 # Purge Dask worker and log directories
 rm -rf $LOGDIR/*
@@ -30,15 +34,17 @@ mkdir -p $WORKER_DIR
 # Purge Dask config directories
 rm -rf ~/.config/dask
 
-# Activate conda environment
-source $CONDA_ENV_PATH
+if [ "${CONDA_DEFAULT_ENV}" != "${CONDA_ENV_NAME}" ] ; then
+  # Activate conda environment
+  source $CONDA_ENV_PATH
+fi
 conda activate $CONDA_ENV_NAME
 
 # Dask/distributed configuration
-export DASK_DISTRIBUTED__COMM__TIMEOUTS__CONNECT="100s"
-export DASK_DISTRIBUTED__COMM__TIMEOUTS__TCP="600s"
-export DASK_DISTRIBUTED__COMM__RETRY__DELAY__MIN="1s"
-export DASK_DISTRIBUTED__COMM__RETRY__DELAY__MAX="60s"
+export DASK_DISTRIBUTED__COMM__TIMEOUTS__CONNECT=${DASK_DISTRIBUTED__COMM__TIMEOUTS__CONNECT:="100s"}
+export DASK_DISTRIBUTED__COMM__TIMEOUTS__TCP=${DASK_DISTRIBUTED__COMM__TIMEOUTS__TCP:="600s"}
+export DASK_DISTRIBUTED__COMM__RETRY__DELAY__MIN=${DASK_DISTRIBUTED__COMM__RETRY__DELAY__MIN:="1s"}
+export DASK_DISTRIBUTED__COMM__RETRY__DELAY__MAX=${DASK_DISTRIBUTED__COMM__RETRY__DELAY__MAX:="60s"}
 
 
 # Setup scheduler

tpcx_bb/cluster_configuration/cluster-startup.sh

@@ -1,25 +1,29 @@
+#!/usr/bin/env bash
+
+set -x
+
 #IB, NVLINK, or TCP
-CLUSTER_MODE=$1
+CLUSTER_MODE=${CLUSTER_MODE:="IB"}
 USERNAME=$(whoami)
 
 MAX_SYSTEM_MEMORY=$(free -m | awk '/^Mem:/{print $2}')M
-DEVICE_MEMORY_LIMIT="25GB"
-POOL_SIZE="30GB"
+DEVICE_MEMORY_LIMIT=${DEVICE_MEMORY_LIMIT:="25GB"}
+POOL_SIZE=${POOL_SIZE:="30GB"}
 
 # Fill in your environment name and conda path on each node
-TPCX_BB_HOME="/home/$USERNAME/shared/tpcx-bb"
-CONDA_ENV_NAME="rapids-tpcx-bb"
-CONDA_ENV_PATH="/home/$USERNAME/conda/etc/profile.d/conda.sh"
+TPCX_BB_HOME=${TPCX_BB_HOME:="/home/$USERNAME/shared/tpcx-bb"}
+CONDA_ENV_NAME=${CONDA_ENV_NAME:="rapids-tpcx-bb"}
+CONDA_ENV_PATH=${CONDA_ENV_PATH:="/home/$USERNAME/conda/etc/profile.d/conda.sh"}
 
 # TODO: Unify interface/IP setting/getting for cluster startup
 # and scheduler file
-INTERFACE="ib0"
+INTERFACE=${INTERFACE:="ib0"}
 
 # TODO: Remove hard-coding of scheduler
-SCHEDULER=$(hostname)
-SCHEDULER_FILE=$TPCX_BB_HOME/tpcx_bb/cluster_configuration/example-cluster-scheduler.json
-LOGDIR="/tmp/tpcx-bb-dask-logs/"
-WORKER_DIR="/tmp/tpcx-bb-dask-workers/"
+SCHEDULER=${SCHEDULER:=$(hostname)}
+SCHEDULER_FILE=${SCHEDULER_FILE:="${TPCX_BB_HOME}/tpcx_bb/cluster_configuration/example-cluster-scheduler.json"}
+LOGDIR=${LOGDIR:="/tmp/tpcx-bb-dask-logs/"}
+WORKER_DIR=${WORKER_DIR:="/tmp/tpcx-bb-dask-workers/"}
 
 # Purge Dask worker and log directories
 rm -rf $LOGDIR/*
@@ -35,10 +39,10 @@ source $CONDA_ENV_PATH
 conda activate $CONDA_ENV_NAME
 
 # Dask/distributed configuration
-export DASK_DISTRIBUTED__COMM__TIMEOUTS__CONNECT="100s"
-export DASK_DISTRIBUTED__COMM__TIMEOUTS__TCP="600s"
-export DASK_DISTRIBUTED__COMM__RETRY__DELAY__MIN="1s"
-export DASK_DISTRIBUTED__COMM__RETRY__DELAY__MAX="60s"
+export DASK_DISTRIBUTED__COMM__TIMEOUTS__CONNECT=${DASK_DISTRIBUTED__COMM__TIMEOUTS__CONNECT:="100s"}
+export DASK_DISTRIBUTED__COMM__TIMEOUTS__TCP=${DASK_DISTRIBUTED__COMM__TIMEOUTS__TCP:="600s"}
+export DASK_DISTRIBUTED__COMM__RETRY__DELAY__MIN=${DASK_DISTRIBUTED__COMM__RETRY__DELAY__MIN:="1s"}
+export DASK_DISTRIBUTED__COMM__RETRY__DELAY__MAX=${DASK_DISTRIBUTED__COMM__RETRY__DELAY__MAX:="60s"}
 
 
 # Setup scheduler
@@ -53,11 +57,10 @@ if [ "$HOSTNAME" = $SCHEDULER ]; then
 fi
 
 
+unset TRANSPORT_ARGS
 # Setup workers
 if [ "$CLUSTER_MODE" = "NVLINK" ]; then
-        dask-cuda-worker --device-memory-limit $DEVICE_MEMORY_LIMIT --local-directory $WORKER_DIR  --rmm-pool-size=$POOL_SIZE --memory-limit=$MAX_SYSTEM_MEMORY --enable-tcp-over-ucx --enable-nvlink  --disable-infiniband --scheduler-file $SCHEDULER_FILE >> $LOGDIR/worker.log 2>&1 &
+    TRANSPORT_ARGS="--enable-tcp-over-ucx --enable-nvlink  --disable-infiniband"
 fi
 
-if [ "$CLUSTER_MODE" = "TCP" ]; then
-    dask-cuda-worker --device-memory-limit $DEVICE_MEMORY_LIMIT --local-directory $WORKER_DIR  --rmm-pool-size=$POOL_SIZE --memory-limit=$MAX_SYSTEM_MEMORY --scheduler-file $SCHEDULER_FILE >> $LOGDIR/worker.log 2>&1 &
-fi
+dask-cuda-worker --device-memory-limit $DEVICE_MEMORY_LIMIT --local-directory $WORKER_DIR  --rmm-pool-size=$POOL_SIZE --memory-limit=$MAX_SYSTEM_MEMORY ${TRANSPORT_ARGS} --scheduler-file $SCHEDULER_FILE >> $LOGDIR/worker.log 2>&1 &

tpcx_bb/queries/load_test/tpcx_bb_load_test.py

@@ -1,11 +1,7 @@
 from xbb_tools.utils import benchmark, tpcxbb_argparser, run_query
 from xbb_tools.readers import build_reader
 import os, subprocess, math, time
-
-
-config["data_dir"] = "/".join(config["data_dir"].rstrip("/").split("/")[:-1])
-
-spark_schema_dir = f"{os.getcwd()}/../../spark_table_schemas/"
+from xbb_tools.config import get_config
 
 # these tables have extra data produced by bigbench dataGen
 refresh_tables = [
@@ -21,15 +17,13 @@
     "web_returns",
     "web_sales",
 ]
-tables = [table.split(".")[0] for table in os.listdir(spark_schema_dir)]
 
-scale = [x for x in config["data_dir"].split("/") if "sf" in x][0]
-part_size = 3
-chunksize = "128 MiB"
+def get_tables( spark_schema_dir ):
+    return [table.split(".")[0] for table in os.listdir(spark_schema_dir)]
 
 # Spark uses different names for column types, and RAPIDS doesn't yet support Decimal types.
-def get_schema(table):
-    with open(f"{spark_schema_dir}{table}.schema") as fp:
+def get_schema(table, schema_dir):
+    with open(os.path.join(schema_dir,f"{table}.schema")) as fp:
         schema = fp.read()
         names = [line.replace(",", "").split()[0] for line in schema.split("\n")]
         types = [
@@ -45,12 +39,9 @@ def get_schema(table):
         return names, types
 
 
-def read_csv_table(table, chunksize="256 MiB"):
-    # build dict of dtypes to use when reading CSV
-    names, types = get_schema(table)
-    dtype = {names[i]: types[i] for i in range(0, len(names))}
-
-    data_dir = config["data_dir"]
+def read_csv_table(table, data_dir, schema_dir, chunksize="256 MiB"):
+    names, types = get_schema(table, schema_dir)
+    dtype=dict(zip(names, types))
     base_path = f"{data_dir}/data/{table}"
     files = os.listdir(base_path)
     # item_marketprices has "audit" files that should be excluded
@@ -99,26 +90,26 @@ def multiplier(unit):
 
 
 # we use size of the CSV data on disk to determine number of Parquet partitions
-def get_size_gb(table):
-    data_dir = config["data_dir"]
-    path = data_dir + "/data/" + table
-    size = subprocess.check_output(["du", "-sh", path]).split()[0].decode("utf-8")
+def get_size_gb(table, data_dir):
+    table_path=os.path.join( data_dir, 'data', table)
+    print( f"Getting size of {table_path}")
+    size = subprocess.check_output(["du", "-sh", table_path]).split()[0].decode("utf-8")
     unit = size[-1]
 
     size = math.ceil(float(size[:-1])) * multiplier(unit)
 
     if table in refresh_tables:
-        path = data_dir + "/data_refresh/" + table
+        table_path = os.path.join(data_dir,'data_refresh',table)
         refresh_size = (
-            subprocess.check_output(["du", "-sh", path]).split()[0].decode("utf-8")
+            subprocess.check_output(["du", "-sh", table_path]).split()[0].decode("utf-8")
         )
         size = size + math.ceil(float(refresh_size[:-1])) * multiplier(refresh_size[-1])
 
     return size
 
 
-def repartition(table, outdir, npartitions=None, chunksize=None, compression="snappy"):
-    size = get_size_gb(table)
+def repartition(table, data_dir, schema_dir, outdir, npartitions=None, chunksize=None, compression="snappy"):
+    size = get_size_gb(table, data_dir)
     if npartitions is None:
         npartitions = max(1, size)
 
@@ -128,38 +119,39 @@ def repartition(table, outdir, npartitions=None, chunksize=None, compression="sn
     # web_clickstreams is particularly memory intensive
     # we sacrifice a bit of speed for stability, converting half at a time
     if table in ["web_clickstreams"]:
-        df = read_csv_table(table, chunksize)
-        half = int(df.npartitions / 2)
-        df.partitions[0:half].repartition(npartitions=int(npartitions / 2)).to_parquet(
+        df = read_csv_table(table, data_dir, schema_dir, chunksize)
+        half = max(1,int(df.npartitions / 2))
+        df.partitions[0:half].repartition(npartitions=max(1,int(npartitions / 2))).to_parquet(
             outdir + table, compression=compression
         )
         print("Completed first half of web_clickstreams..")
-        df.partitions[half:].repartition(npartitions=int(npartitions / 2)).to_parquet(
+        df.partitions[half:].repartition(npartitions=max(1,int(npartitions / 2))).to_parquet(
             outdir + table, compression=compression
         )
 
     else:
-        read_csv_table(table, chunksize).repartition(
+        read_csv_table(table, data_dir, schema_dir, chunksize).repartition(
             npartitions=npartitions
         ).to_parquet(outdir + table, compression=compression)
 
 
 def main(client, config):
     # location you want to write Parquet versions of the table data
-    data_dir = "/".join(config["data_dir"].split("/")[:-1])
-    outdir = f"{data_dir}/parquet_{part_size}gb/"
-
-    total = 0
-    for table in tables:
-        size_gb = get_size_gb(table)
+    data_dir = config.get("data_dir",'.')
+    outdir = f"{config.get('output_dir',data_dir)}/parquet_{config.get('partitions',3)}gb/"
+    schema_dir= config.get('spark_schema_dir',
+                           os.path.join(os.getcwd(),'..','..','spark_table_schemas'))
+    began = time.time()
+    for table in get_tables(schema_dir):
+        size_gb = get_size_gb(table, data_dir)
         # product_reviews has lengthy strings which exceed cudf's max number of characters per column
         # we use smaller partitions to avoid overflowing this character limit
         if table == "product_reviews":
             npartitions = max(1, int(size_gb / 1))
         else:
-            npartitions = max(1, int(size_gb / part_size))
-        repartition(table, outdir, npartitions, chunksize, compression="snappy")
-    print(f"{chunksize} took {total}s")
+            npartitions = max(1, int(size_gb / config.get('partitions',3)))
+        repartition(table, data_dir, schema_dir, outdir, npartitions, config.get('chunk_size',"128 MiB"), compression="snappy")
+    print(f"{config.get('chunk_size','128 MiB')} took {time.time() - began}s")
     return cudf.DataFrame()
 
 
@@ -168,6 +160,6 @@ def main(client, config):
     import cudf
     import dask_cudf
 
-    config = tpcxbb_argparser()
-    client, bc = attach_to_cluster(config)
-    run_query(config=config, client=client, query_func=main)
+    conf = tpcxbb_argparser()
+    client, bc = attach_to_cluster(conf)
+    run_query(config=conf, client=client, query_func=main)