diff --git a/llama_stack/distribution/start_container.sh b/llama_stack/distribution/start_container.sh
index fe1b5051f..b9ec9a23d 100755
--- a/llama_stack/distribution/start_container.sh
+++ b/llama_stack/distribution/start_container.sh
@@ -10,6 +10,8 @@ DOCKER_BINARY=${DOCKER_BINARY:-docker}
 DOCKER_OPTS=${DOCKER_OPTS:-}
 LLAMA_CHECKPOINT_DIR=${LLAMA_CHECKPOINT_DIR:-}
 LLAMA_STACK_DIR=${LLAMA_STACK_DIR:-}
+TEST_PYPI_VERSION=${TEST_PYPI_VERSION:-}
+PYPI_VERSION=${PYPI_VERSION:-}
 
 set -euo pipefail
 
@@ -54,6 +56,12 @@ if [ -n "$LLAMA_CHECKPOINT_DIR" ]; then
   DOCKER_OPTS="$DOCKER_OPTS --gpus=all"
 fi
 
+if [ -n "$PYPI_VERSION" ]; then
+  docker_image="$docker_image-$PYPI_VERSION"
+elif [ -n "$TEST_PYPI_VERSION" ]; then
+  docker_image="$docker_image-test-$TEST_PYPI_VERSION"
+fi
+
 $DOCKER_BINARY run $DOCKER_OPTS -it \
   -p $port:$port \
   -v "$yaml_config:/app/config.yaml" \