Update.

Signed-off-by: Govind Kamat <govkamat@amazon.com>

docker/Dockerfile

@@ -2,18 +2,24 @@
 # Install OpenSearch Benchmark to build a Docker image #
 ########################################################
 
+#
+# Stage 1: build packages and compile where needed
+#
 FROM python:3.11.2-slim AS build-stage
 WORKDIR /opensearch-benchmark
-ENV BENCHMARK_RUNNING_IN_DOCKER=True
 
 RUN apt-get -y update && \
     apt-get install -y curl git gcc pbzip2 pigz make jq && \
-    apt-get -y upgrade && \
-    rm -rf /var/lib/apt/lists/*
+    apt-get -y upgrade
 
 COPY . opensearch-benchmark
+
 RUN cd opensearch-benchmark; make wheel
 
+
+#
+# Stage 2: create image
+#
 FROM python:3.11.2-slim AS image-stage
 WORKDIR /opensearch-benchmark
 ENV BENCHMARK_RUNNING_IN_DOCKER=True
@@ -24,8 +30,16 @@ RUN groupadd --gid 1000 opensearch-benchmark && \
 RUN mkdir -p /opensearch-benchmark/.benchmark && \
     chown -R 1000:0 /opensearch-benchmark/.benchmark
 
-COPY --from=build-stage /opensearch-benchmark/opensearch-benchmark/yappi-*.whl /opensearch-benchmark/opensearch-benchmark/opensearch_benchmark-*.whl .
-RUN set -ex; apt-get -y update; apt-get install -y git pbzip2; PIP_ONLY_BINARY=h5py pip install yappi-*.whl opensearch_benchmark-*.whl; rm *.whl
+COPY --from=build-stage /opensearch-benchmark/opensearch-benchmark/yappi-*.whl /opensearch-benchmark/opensearch-benchmark/opensearch_benchmark-*.whl ./
+
+# There is no binary package currently available for yappi on ARM.
+RUN set -ex; \
+    apt-get -y update; \
+    apt-get install -y git pbzip2; \
+    apt-get -y upgrade; \
+    rm -rf /var/lib/apt/lists/*; \
+    PIP_ONLY_BINARY=h5py pip install yappi-*.whl opensearch_benchmark-*.whl; \
+    rm *.whl
 
 USER 1000