Allow setting environment variables from llama stack run and fix ollama

llama_stack/distribution/build_container.sh

@@ -146,6 +146,8 @@ fi
 # Set version tag based on PyPI version
 if [ -n "$TEST_PYPI_VERSION" ]; then
   version_tag="test-$TEST_PYPI_VERSION"
+elif [[ -n "$LLAMA_STACK_DIR" || -n "$LLAMA_MODELS_DIR" ]]; then
+  version_tag="dev"
 else
   URL="https://pypi.org/pypi/llama-stack/json"
   version_tag=$(curl -s $URL | jq -r '.info.version')

llama_stack/distribution/start_conda_env.sh

@@ -33,10 +33,33 @@ shift
 port="$1"
 shift
 
+# Process environment variables from --env arguments
+env_vars=""
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --env)
+
+            if [[ -n "$2" ]]; then
+                # collect environment variables so we can set them after activating the conda env
+                env_vars="$env_vars $2"
+                shift 2
+            else
+                echo -e "${RED}Error: --env requires a KEY=VALUE argument${NC}" >&2
+                exit 1
+            fi
+            ;;
+        *)
+            shift
+            ;;
+    esac
+done
+
 eval "$(conda shell.bash hook)"
 conda deactivate && conda activate "$env_name"
 
-$CONDA_PREFIX/bin/python \
+set -x
+$env_vars \
+  $CONDA_PREFIX/bin/python \
   -m llama_stack.distribution.server.server \
   --yaml_config "$yaml_config" \
   --port "$port" "$@"

llama_stack/distribution/start_container.sh

@@ -31,7 +31,7 @@ if [ $# -lt 3 ]; then
 fi
 
 build_name="$1"
-docker_image="distribution-$build_name"
+docker_image="localhost/distribution-$build_name"
 shift
 
 yaml_config="$1"
@@ -40,6 +40,26 @@ shift
 port="$1"
 shift
 
+# Process environment variables from --env arguments
+env_vars=""
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        --env)
+            echo "env = $2"
+            if [[ -n "$2" ]]; then
+                env_vars="$env_vars -e $2"
+                shift 2
+            else
+                echo -e "${RED}Error: --env requires a KEY=VALUE argument${NC}" >&2
+                exit 1
+            fi
+            ;;
+        *)
+            shift
+            ;;
+    esac
+done
+
 set -x
 
 if command -v selinuxenabled &> /dev/null && selinuxenabled; then
@@ -59,15 +79,18 @@ fi
 version_tag="latest"
 if [ -n "$PYPI_VERSION" ]; then
   version_tag="$PYPI_VERSION"
+elif [ -n "$LLAMA_STACK_DIR" ]; then
+  version_tag="dev"
 elif [ -n "$TEST_PYPI_VERSION" ]; then
   version_tag="test-$TEST_PYPI_VERSION"
 fi
 
 $DOCKER_BINARY run $DOCKER_OPTS -it \
   -p $port:$port \
+  $env_vars \
   -v "$yaml_config:/app/config.yaml" \
   $mounts \
   $docker_image:$version_tag \
   python -m llama_stack.distribution.server.server \
   --yaml_config /app/config.yaml \
-  --port $port "$@"
+  --port "$port"