diff --git a/docker/sandbox-bundled/Dockerfile b/docker/sandbox-bundled/Dockerfile
index 92e8fa2d46..0c7a672f9d 100644
--- a/docker/sandbox-bundled/Dockerfile
+++ b/docker/sandbox-bundled/Dockerfile
@@ -1,49 +1,59 @@
 # syntax=docker/dockerfile:1.4-labs
+
+# Stage 1: Builder stage for preloading images
 FROM --platform=${BUILDPLATFORM} mgoltzsche/podman:minimal AS builder
 
 ARG TARGETARCH
-ENV TARGETARCH "${TARGETARCH}"
+ENV TARGETARCH="${TARGETARCH}"
 
 WORKDIR /build
 
+
 COPY images/manifest.txt images/preload ./
 RUN --security=insecure ./preload manifest.txt
 
 
+# Stage 2: Bootstrap stage to compile Go binary
 FROM --platform=${BUILDPLATFORM} golang:1.22-bullseye AS bootstrap
 
 ARG TARGETARCH
-ENV CGO_ENABLED 0
-ENV GOARCH "${TARGETARCH}"
-ENV GOOS linux
+ENV CGO_ENABLED=0
+ENV GOARCH="${TARGETARCH}"
+ENV GOOS=linux
 
 WORKDIR /flyteorg/build
+
+# Separate dependency caching step for Go modules
 COPY bootstrap/go.mod bootstrap/go.sum ./
-RUN go mod download
-COPY bootstrap/ ./
 RUN --mount=type=cache,target=/root/.cache/go-build --mount=type=cache,target=/root/go/pkg/mod \
-    go build -o dist/flyte-sandbox-bootstrap cmd/bootstrap/main.go
+    go mod download
+
+# Build the bootstrap binary
+COPY bootstrap/ ./
+RUN go build -o dist/flyte-sandbox-bootstrap cmd/bootstrap/main.go
 
 
 FROM rancher/k3s:v1.29.0-k3s1
 
 ARG TARGETARCH
-
+ENV TARGETARCH="${TARGETARCH}"
 ARG FLYTE_SANDBOX_VERSION
-ENV FLYTE_SANDBOX_VERSION "${FLYTE_SANDBOX_VERSION}"
+ENV FLYTE_SANDBOX_VERSION="${FLYTE_SANDBOX_VERSION}"
 
+# Copy preloaded images and manifests for k3s
 COPY --from=builder /build/images/ /var/lib/rancher/k3s/agent/images/
 COPY images/tar/${TARGETARCH}/ /var/lib/rancher/k3s/agent/images/
 COPY manifests/ /var/lib/rancher/k3s/server/manifests-staging/
 COPY bin/ /bin/
 
-# Install bootstrap
+
 COPY --from=bootstrap /flyteorg/build/dist/flyte-sandbox-bootstrap /bin/
 
+# Volume for Flyte storage
 VOLUME /var/lib/flyte/storage
 
-# Set environment variable for picking up additional CA certificates
-ENV SSL_CERT_DIR /var/lib/flyte/config/ca-certificates
+
+ENV SSL_CERT_DIR=/var/lib/flyte/config/ca-certificates
 
 ENTRYPOINT [ "/bin/k3d-entrypoint.sh" ]
 CMD [ "server", "--disable=traefik", "--disable=servicelb" ]
diff --git a/docker/sandbox/Dockerfile b/docker/sandbox/Dockerfile
index 40c3e83e83..be76841660 100644
--- a/docker/sandbox/Dockerfile
+++ b/docker/sandbox/Dockerfile
@@ -1,73 +1,69 @@
-FROM alpine:3.13.5 AS base
+# syntax=docker/dockerfile:1.4-labs
 
-# Install dependencies
-RUN apk add --no-cache openssl
+# Stage 1: Base setup with Alpine
+FROM alpine:3.20.3 AS base
 
-# Make directory to store artifacts
+# Install necessary dependencies
+RUN apk add --no-cache openssl bash git make tini curl jq
+
+# Set up directories for Flyte artifacts
 RUN mkdir -p /flyteorg/bin /flyteorg/share
 
-# Install k3s
-ARG K3S_VERSION="v1.21.1%2Bk3s1"
+# Install k3s with version and architecture specification
+ARG K3S_VERSION="v1.21.1+k3s1"
 ARG TARGETARCH
 
 RUN case $TARGETARCH in \
     amd64) export SUFFIX=;; \
-    arm64) export SUFFIX=-arm64;; \
-    aarch64)  export SUFFIX=-arm64;; \
-    # TODO: Check if we need to add case fail
+    arm64 | aarch64) export SUFFIX=-arm64;; \
+    *) echo "Unsupported architecture"; exit 1;; \
     esac; \
     wget -q -O /flyteorg/bin/k3s https://github.com/k3s-io/k3s/releases/download/${K3S_VERSION}/k3s${SUFFIX} \
     && chmod +x /flyteorg/bin/k3s
 
-# Install Helm
+# Install Helm 
 ARG HELM_VERSION="v3.6.3"
-
 RUN wget -q -O /flyteorg/bin/get_helm.sh https://raw.githubusercontent.com/helm/helm/${HELM_VERSION}/scripts/get-helm-3 && \
     chmod 700 /flyteorg/bin/get_helm.sh && \
     sh /flyteorg/bin/get_helm.sh --version ${HELM_VERSION} && \
     mv /usr/local/bin/helm /flyteorg/bin/helm && \
     rm /flyteorg/bin/get_helm.sh
 
-
 # Install flytectl
 RUN wget -q -O - https://raw.githubusercontent.com/flyteorg/flytectl/master/install.sh | BINDIR=/flyteorg/bin sh -s
 
-# Copy flyte chart
+# Copy Flyte Helm charts and scripts
 COPY charts/flyte/ /flyteorg/share/flyte
 COPY charts/flyte-core/ /flyteorg/share/flyte-core
 COPY charts/flyte-deps/ /flyteorg/share/flyte-deps
-
-# Copy scripts
 COPY docker/sandbox/kubectl docker/sandbox/cgroup-v2-hack.sh docker/sandbox/wait-for-flyte.sh /flyteorg/bin/
 
-FROM docker:20.10.14-dind-alpine3.15 AS dind
 
-# Install dependencies
-RUN apk add --no-cache bash git make tini curl jq
+# Stage 2: Docker-in-Docker setup with dependencies
+FROM docker:20.10.25-dind-alpine3.20 AS dind
 
-# Copy artifacts from base
+# Copy artifacts from base stage
 COPY --from=base /flyteorg/ /flyteorg/
 
-# Copy entrypoints
-COPY docker/sandbox/flyte-entrypoint-default.sh /flyteorg/bin/flyte-entrypoint.sh
-
+# Set environment variables for Flyte version and test type
 ARG FLYTE_VERSION="latest"
 ENV FLYTE_VERSION="${FLYTE_VERSION}"
-
 ARG FLYTE_TEST="release"
 ENV FLYTE_TEST="${FLYTE_TEST}"
 
-# Update PATH variable
+# Update PATH to include Flyte binaries
 ENV PATH="/flyteorg/bin:${PATH}"
 
-# Declare volumes for k3s
+# Declare volumes for k3s to persist data across containers
 VOLUME /var/lib/kubelet
 VOLUME /var/lib/rancher/k3s
 VOLUME /var/lib/cni
 VOLUME /var/log
 
-# Expose Flyte ports
+# Expose Flyte service ports
 EXPOSE 30081 30082 30084 30088
 
+# Copy entrypoint script and set it as entrypoint
+COPY docker/sandbox/flyte-entrypoint-default.sh /flyteorg/bin/flyte-entrypoint.sh
+ENTRYPOINT ["tini", "--", "/flyteorg/bin/flyte-entrypoint.sh"]
 
-ENTRYPOINT ["tini", "flyte-entrypoint.sh"]