From 00fa7da12535b77625f7140c3f377b64b4d408fa Mon Sep 17 00:00:00 2001
From: Scott Todd <scott.todd0@gmail.com>
Date: Thu, 12 Dec 2024 09:33:48 -0800
Subject: [PATCH] Enable tokenizers in shortfin packages on Linux x86_64.

---
 shortfin/build_tools/build_linux_package.sh | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/shortfin/build_tools/build_linux_package.sh b/shortfin/build_tools/build_linux_package.sh
index 91b944e51..b812f3d89 100755
--- a/shortfin/build_tools/build_linux_package.sh
+++ b/shortfin/build_tools/build_linux_package.sh
@@ -37,13 +37,22 @@ REPO_ROOT="$(cd "$THIS_DIR"/../../ && pwd)"
 SCRIPT_NAME="$(basename $0)"
 ARCH="$(uname -m)"
 
-# Note: we can switch to https://github.com/nod-ai/base-docker-images as needed for extra deps.
-MANYLINUX_DOCKER_IMAGE="${MANYLINUX_DOCKER_IMAGE:-quay.io/pypa/manylinux_2_28_${ARCH}:latest}"
 PYTHON_VERSIONS="${OVERRIDE_PYTHON_VERSIONS:-cp311-cp311 cp312-cp312 cp313-cp313}"
 OUTPUT_DIR="${OUTPUT_DIR:-${THIS_DIR}/wheelhouse}"
 CACHE_DIR="${CACHE_DIR:-}"
 SHORTFIN_ENABLE_TRACING="${SHORTFIN_ENABLE_TRACING:-ON}"
 
+# Our x86_64 dockerfile contains dependencies needed to build the tokenizer
+# library (rust), so enable tokenizers by default on x86_64.
+# TODO: publish a multi-platform manylinux image and include more deps in all platforms (rust, ccache, etc.)
+if [[ "${ARCH}" == "x86_64" ]]; then
+  MANYLINUX_DOCKER_IMAGE="${MANYLINUX_DOCKER_IMAGE:-ghcr.io/nod-ai/manylinux_x86_64@sha256:4acf83343706d1e37252d6001ded3c97a73bc38620580f855b4e65e35ddc5681}"
+  SHORTFIN_ENABLE_TOKENIZERS="${SHORTFIN_ENABLE_TOKENIZERS:-ON}"
+else
+  MANYLINUX_DOCKER_IMAGE="${MANYLINUX_DOCKER_IMAGE:-quay.io/pypa/manylinux_2_28_${ARCH}:latest}"
+  SHORTFIN_ENABLE_TOKENIZERS="${SHORTFIN_ENABLE_TOKENIZERS:-OFF}"
+fi
+
 function run_on_host() {
   echo "Running on host"
   echo "Launching docker image ${MANYLINUX_DOCKER_IMAGE}"
@@ -69,6 +78,7 @@ function run_on_host() {
     -e "OVERRIDE_PYTHON_VERSIONS=${PYTHON_VERSIONS}" \
     -e "OUTPUT_DIR=${OUTPUT_DIR}" \
     -e "SHORTFIN_ENABLE_TRACING=${SHORTFIN_ENABLE_TRACING}" \
+    -e "SHORTFIN_ENABLE_TOKENIZERS=${SHORTFIN_ENABLE_TOKENIZERS}" \
     ${extra_args} \
     "${MANYLINUX_DOCKER_IMAGE}" \
     -- ${THIS_DIR}/${SCRIPT_NAME}