From a57b7108c8e90bcc9894cf71809803466db29faa Mon Sep 17 00:00:00 2001
From: Vivian Wu <vivianrwu@google.com>
Date: Fri, 10 Jan 2025 20:49:25 +0000
Subject: [PATCH] add environment variable to change jetstream version

---
 inference/maxengine_server/Dockerfile | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/inference/maxengine_server/Dockerfile b/inference/maxengine_server/Dockerfile
index b94ea499a..7e5059768 100644
--- a/inference/maxengine_server/Dockerfile
+++ b/inference/maxengine_server/Dockerfile
@@ -19,6 +19,7 @@ FROM ubuntu:22.04
 
 ENV DEBIAN_FRONTEND=noninteractive
 ENV MAXTEXT_VERSION=main
+ENV JETSTREAM_VERSION=main
 
 RUN apt -y update && apt install -y --no-install-recommends \
     ca-certificates \
@@ -37,6 +38,7 @@ git checkout ${MAXTEXT_VERSION} && \
 bash setup.sh
 
 RUN cd /JetStream && \
+git checkout ${JETSTREAM_VERSION} && \
 pip install -e .
 
 COPY maxengine_server_entrypoint.sh /usr/bin/