From ecdcfb28ca107665b7653e237064ee31883c4544 Mon Sep 17 00:00:00 2001
From: r3v5 <milleryan2003@gmail.com>
Date: Mon, 21 Jul 2025 10:26:09 +0100
Subject: [PATCH] feat: allow Distribution container images running rootlessly
 in Llama Stack

---
 llama_stack/distribution/build_container.sh | 29 ++++++++++++++++-----
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/llama_stack/distribution/build_container.sh b/llama_stack/distribution/build_container.sh
index 6985c1cd0..2e85f8764 100755
--- a/llama_stack/distribution/build_container.sh
+++ b/llama_stack/distribution/build_container.sh
@@ -20,7 +20,7 @@ UV_HTTP_TIMEOUT=${UV_HTTP_TIMEOUT:-500}
 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"}
+MOUNT_CACHE=${MOUNT_CACHE:-"--mount=type=cache,id=llama-stack-cache,target=~/.cache"}
 
 # Path to the run.yaml file in the container
 RUN_CONFIG_PATH=/app/run.yaml
@@ -259,6 +259,27 @@ fi
 RUN pip uninstall -y uv
 EOF
 
+# Add non-root user setup before entrypoint
+add_to_container << EOF
+
+# Create group with GID 1001 and user with UID 1001
+RUN groupadd -g 1001 1001 && useradd -u 1001 -g 1001 -M 1001
+
+# Create necessary directories with appropriate permissions for UID 1001
+# Set SELinux context to allow access by non-root user (if SELinux is supported)
+RUN mkdir -p /.llama /.cache && chown -R 1001:1001 /.llama /.cache && chmod -R 775 /.llama /.cache && chmod -R g+w /app && \
+    (command -v chcon >/dev/null 2>&1 && touch /tmp/selinux_test && chcon -t container_file_t /tmp/selinux_test 2>/dev/null && chcon -R -t container_file_t /.llama /.cache 2>/dev/null && rm -f /tmp/selinux_test || true)
+
+# Set the Llama Stack config directory environment variable to use /.llama
+ENV LLAMA_STACK_CONFIG_DIR=/.llama
+
+# This prevents dual storage while keeping /app as working directory for CI compatibility
+ENV HOME=/
+
+# Switch to non-root user (UID 1001 directly)
+USER 1001
+EOF
+
 # If a run config is provided, we use the --config flag
 if [[ -n "$run_config" ]]; then
   add_to_container << EOF
@@ -271,12 +292,6 @@ ENTRYPOINT ["python", "-m", "llama_stack.distribution.server.server", "--templat
 EOF
 fi
 
-# Add other require item commands genearic to all containers
-add_to_container << EOF
-
-RUN mkdir -p /.llama /.cache && chmod -R g+rw /app /.llama /.cache
-EOF
-
 printf "Containerfile created successfully in %s/Containerfile\n\n" "$TEMP_DIR"
 cat "$TEMP_DIR"/Containerfile
 printf "\n"