From 340448e0aa0540f824ce0476505594d55f1088e6 Mon Sep 17 00:00:00 2001
From: Derek Higgins <derekh@redhat.com>
Date: Wed, 23 Jul 2025 00:51:52 +0100
Subject: [PATCH] fix: optimize container build by enabling uv cache (#2855)

- Remove --no-cache flags from uv pip install commands to enable caching
- Mount host uv cache directory to container for persistent caching
- Set UV_LINK_MODE=copy to prevent uv using hardlinks
- When building the starter image
o Build time reduced from ~4:45 to ~3:05 on subsequent builds
(environment specific)
  o Eliminates re-downloading of 3G+ of data on each build
  o Cache size: ~6.2G (when building starter image)

Fixes excessive data downloads during distro container builds.

Signed-off-by: Derek Higgins <derekh@redhat.com>
---
 llama_stack/distribution/build_container.sh | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/llama_stack/distribution/build_container.sh b/llama_stack/distribution/build_container.sh
index 6e794b36f..74776dd7d 100755
--- a/llama_stack/distribution/build_container.sh
+++ b/llama_stack/distribution/build_container.sh
@@ -19,6 +19,9 @@ UV_HTTP_TIMEOUT=${UV_HTTP_TIMEOUT:-500}
 # mounting is not supported by docker buildx, so we use COPY instead
 USE_COPY_NOT_MOUNT=${USE_COPY_NOT_MOUNT:-}
 
+# Mount command for cache container .cache, can be overridden by the user if needed
+MOUNT_CACHE=${MOUNT_CACHE:-"--mount=type=cache,id=llama-stack-cache,target=/root/.cache"}
+
 # Path to the run.yaml file in the container
 RUN_CONFIG_PATH=/app/run.yaml
 
@@ -125,11 +128,16 @@ RUN pip install uv
 EOF
 fi
 
+# Set the link mode to copy so that uv doesn't attempt to symlink to the cache directory
+add_to_container << EOF
+ENV UV_LINK_MODE=copy
+EOF
+
 # Add pip dependencies first since llama-stack is what will change most often
 # so we can reuse layers.
 if [ -n "$pip_dependencies" ]; then
   add_to_container << EOF
-RUN uv pip install --no-cache $pip_dependencies
+RUN $MOUNT_CACHE uv pip install $pip_dependencies
 EOF
 fi
 
@@ -137,7 +145,7 @@ if [ -n "$special_pip_deps" ]; then
   IFS='#' read -ra parts <<<"$special_pip_deps"
   for part in "${parts[@]}"; do
     add_to_container <<EOF
-RUN uv pip install --no-cache $part
+RUN $MOUNT_CACHE uv pip install $part
 EOF
   done
 fi
@@ -207,7 +215,7 @@ COPY $dir $mount_point
 EOF
   fi
   add_to_container << EOF
-RUN uv pip install --no-cache -e $mount_point
+RUN $MOUNT_CACHE uv pip install -e $mount_point
 EOF
 }
 
@@ -222,10 +230,10 @@ else
   if [ -n "$TEST_PYPI_VERSION" ]; then
     # these packages are damaged in test-pypi, so install them first
     add_to_container << EOF
-RUN uv pip install fastapi libcst
+RUN $MOUNT_CACHE uv pip install fastapi libcst
 EOF
     add_to_container << EOF
-RUN uv pip install --no-cache --extra-index-url https://test.pypi.org/simple/ \
+RUN $MOUNT_CACHE uv pip install --extra-index-url https://test.pypi.org/simple/ \
   --index-strategy unsafe-best-match \
   llama-stack==$TEST_PYPI_VERSION
 
@@ -237,7 +245,7 @@ EOF
       SPEC_VERSION="llama-stack"
     fi
     add_to_container << EOF
-RUN uv pip install --no-cache $SPEC_VERSION
+RUN $MOUNT_CACHE uv pip install $SPEC_VERSION
 EOF
   fi
 fi