Dockerfile for benchmarking Bert Base Uncased Model in GCP (OpenMP disabled single thread)

Dockerfile

# Install TF entreprise
FROM gcr.io/deeplearning-platform-release/tf2-cpu.2-5:latest
WORKDIR /workspace

# OpenMP disabled in versions >=  1.7.0 (https://github.com/microsoft/onnxruntime/releases/tag/v1.7.0)
RUN pip install --no-cache-dir --upgrade pip && \
    pip install --upgrade torch==1.8.1+cpu -f https://download.pytorch.org/whl/torch_stable.html && \
    pip install coloredlogs sympy onnx tf2onnx onnxruntime==1.7.0 transformers==4.6.0

# Service account credentials are required to store benchmark results in a GCS bucket. 
# Requirements before building this image.
# 1. Create a gcs bucket named gs://bert-inference-results
# 2. Create a service account and grant it access to the GCS bucket.
COPY ./<SERVICE_ACCOUNT_KEY_FILE>.json /workspace/<SERVICE_ACCOUNT_KEY_FILE>.json
RUN gcloud auth activate-service-account --key-file=/workspace/<SERVICE_ACCOUNT_KEY_FILE>.json

COPY ./run_single_thread_benchmark.sh /workspace/
COPY ./bert-base-uncased /workspace/bert-base-uncased

# Zip the onnxruntime customized benchmark scripts from https://github.com/mwakaba2/onnxruntime/tree/benchmark-bert-base-uncased
# Add the zipped file to workspace
ADD onnxruntime_benchmark.tar.gz /workspace

# Setting to make sure the script is using only a single thread. 
ENV OMP_NUM_THREADS=1
ENV OMP_WAIT_POLICY='PASSIVE'

CMD [ "./run_single_thread_benchmark.sh" ]

run_single_thread_benchmark.sh can be found here