Add loadgen in dev image (#109)

* build loadgen in docker
AI-Hypercomputer · Jul 2, 2024 · 69ce8a2 · 69ce8a2
1 parent 46e444c
commit 69ce8a2
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/jetstream/tools/proxy_dev/base.Dockerfile b/jetstream/tools/proxy_dev/base.Dockerfile
@@ -5,7 +5,7 @@ FROM ubuntu:22.04
 
 ENV DEBIAN_FRONTEND=noninteractive
 
-RUN apt -y update && apt install -y --no-install-recommends apt-transport-https ca-certificates gnupg git python3.10 python3-pip curl
+RUN apt -y update && apt install -y --no-install-recommends apt-transport-https ca-certificates gnupg git python3.10 python3-pip curl nano vim
 
 RUN update-alternatives --install     /usr/bin/python3 python3 /usr/bin/python3.10 1
 RUN echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | gpg --dearmor -o /usr/share/keyrings/cloud.google.gpg && apt-get update -y && apt-get install google-cloud-sdk -y

diff --git a/jetstream/tools/proxy_dev/dev.Dockerfile b/jetstream/tools/proxy_dev/dev.Dockerfile
@@ -14,4 +14,16 @@ COPY  maxtext ./maxtext
 
 RUN pip install ./JetStream
 
+COPY inference_mlperf4.1 ./inference_mlperf4.1
+RUN apt-get -y install python3-dev && apt-get -y install build-essential
+RUN pip install ./inference_mlperf4.1/loadgen
+RUN pip install \
+    transformers==4.31.0 \
+    nltk==3.8.1 \
+    evaluate==0.4.0 \
+    absl-py==1.4.0 \
+    rouge-score==0.1.2 \
+    sentencepiece==0.1.99 \
+    accelerate==0.21.0
+
 ENTRYPOINT ["bash"]