From 99181d4fee3b20c31073259eab40c7f32c32a847 Mon Sep 17 00:00:00 2001
From: Sax Authors <no-reply@google.com>
Date: Mon, 1 Apr 2024 15:59:51 -0700
Subject: [PATCH] Simplifying using volume mount instead of GCSBucket in
 singlehost case. Changed SAX_CELL default to /sax/cell

PiperOrigin-RevId: 620978624
Change-Id: I445693188c391d6de70c5d916c2cb00c850e1394
---
 saxml/tools/docker/Dockerfile.admin | 11 ++++++++---
 saxml/tools/docker/Dockerfile.model |  8 ++++++--
 2 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/saxml/tools/docker/Dockerfile.admin b/saxml/tools/docker/Dockerfile.admin
index 14aaf04d..3ec350f9 100644
--- a/saxml/tools/docker/Dockerfile.admin
+++ b/saxml/tools/docker/Dockerfile.admin
@@ -37,15 +37,20 @@ COPY --from=build_image \
     /saxml/bazel-bin/saxml/bin/saxutil_/saxutil \
     /usr/bin/saxutil
 
-# Use "/sax/test" as default sax cell.
-ENV SAX_CELL=/sax/test
+# Use "/sax/cell" as default sax cell.
+ENV SAX_CELL=/sax/cell
 # Use "10000" as default port for admin server gRPC endpoint.
 ENV PORT=10000
 ENV SAX_ROOT=""
 ENV FS_ROOT=""
 ENV SAX_GCS_BUCKET=""
+ENV SAX_DIR=""
 # Create a script that configures and runs the admin server.
 RUN echo '#!/bin/bash\n\
+if [[ -n "$SAX_DIR" ]]; then\n\
+	SAX_ROOT=${SAX_DIR}/sax-root \n\
+	FS_ROOT=${SAX_DIR}/sax-fs-root \n\
+fi\n\
 if [[ -n "$GSBUCKET" ]]; then\n\
 	SAX_GCS_BUCKET=${GSBUCKET}\n\
 fi\n\
@@ -72,4 +77,4 @@ admin_server \
   --alsologtostderr\n' \
 > /usr/bin/admin_server_entrypoint.sh && chmod +x /usr/bin/admin_server_entrypoint.sh
 
-ENTRYPOINT ["/usr/bin/admin_server_entrypoint.sh"]
+ENTRYPOINT ["/usr/bin/admin_server_entrypoint.sh"]
\ No newline at end of file
diff --git a/saxml/tools/docker/Dockerfile.model b/saxml/tools/docker/Dockerfile.model
index 76fbc3a5..5bba27da 100644
--- a/saxml/tools/docker/Dockerfile.model
+++ b/saxml/tools/docker/Dockerfile.model
@@ -58,14 +58,18 @@ COPY --from=build_image \
     /usr/bin
 
 ENV SAX_ROOT=""
+ENV SAX_DIR=""
 
 RUN echo '#!/bin/bash\n\
+if [[ -n "$SAX_DIR" ]];then \n\
+	SAX_ROOT=${LOCAL_DIR}/sax-root \n\
+fi\n\
 if [[ -z "$SAX_ROOT"  ]]; then\n\
 	if [[ -n "$SAX_GCS_BUCKET" && -n "$SAX_REPLICA_NAME" ]]; then\n\
 		SAX_ROOT=gs://${SAX_GCS_BUCKET}/${SAX_REPLICA_NAME}/sax-root \n\
 	fi\n\
 fi\n\
-server --alsologtostderr --sax_cell=/sax/test $@\n' \
+server --alsologtostderr --sax_cell=/sax/cell $@\n' \
 > /usr/bin/model_server_entrypoint.sh && chmod +x /usr/bin/model_server_entrypoint.sh
 
-ENTRYPOINT ["/usr/bin/model_server_entrypoint.sh"]
+ENTRYPOINT ["/usr/bin/model_server_entrypoint.sh"]
\ No newline at end of file