perf(docker): optimize Dockerfile build cache with BuildKit mount cache

docker/Dockerfile.cuda

@@ -12,17 +12,16 @@ FROM ${BASE_CUDA_DEV_CONTAINER} AS build
 ARG RUST_TOOLCHAIN=stable
 
 ENV DEBIAN_FRONTEND=noninteractive
-RUN apt-get update && \
+RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
+    --mount=type=cache,target=/var/lib/apt/lists,sharing=locked \
+    apt-get update && \
     apt-get install -y --no-install-recommends \
         curl \
         pkg-config \
         libssl-dev \
         protobuf-compiler \
         git \
-        cmake \
-        && \
-    apt-get clean && \
-    rm -rf /var/lib/apt/lists/*
+        cmake
 
 # setup rust.
 RUN curl https://sh.rustup.rs -sSf | bash -s -- --default-toolchain ${RUST_TOOLCHAIN} -y
@@ -48,17 +47,16 @@ RUN cp /opt/tabby/bin/tabby /opt/tabby/bin/tabby-cpu
 
 FROM ${BASE_CUDA_RUN_CONTAINER} AS runtime
 
-RUN apt-get update && \
+RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
+    --mount=type=cache,target=/var/lib/apt/lists,sharing=locked \
+    apt-get update && \
     apt-get install -y --no-install-recommends \
         git   \
         curl  \
         unzip \
         openssh-client \
         ca-certificates \
-        libgomp1 \
-        && \
-    apt-get clean && \
-    rm -rf /var/lib/apt/lists/*
+        libgomp1
 
 # Install katana
 RUN curl -L https://github.com/projectdiscovery/katana/releases/download/v1.1.2/katana_1.1.2_linux_amd64.zip -o katana.zip \

docker/Dockerfile.rocm

@@ -1,6 +1,6 @@
 ARG UBUNTU_VERSION=22.04
 # This needs to generally match the container host's environment.
-ARG ROCM_VERSION=5.7.1
+ARG ROCM_VERSION=6.1.2
 # Target the ROCM build image
 ARG BASE_ROCM_DEV_CONTAINER=rocm/dev-ubuntu-${UBUNTU_VERSION}:${ROCM_VERSION}-complete
 # Target the ROCM runtime image
@@ -12,17 +12,16 @@ FROM ${BASE_ROCM_DEV_CONTAINER} AS build
 ARG RUST_TOOLCHAIN=stable
 
 ENV DEBIAN_FRONTEND=noninteractive
-RUN apt-get update && \
+RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
+    --mount=type=cache,target=/var/lib/apt/lists,sharing=locked \
+    apt-get update && \
     apt-get install -y --no-install-recommends \
-    curl \
-    pkg-config \
-    libssl-dev \
-    protobuf-compiler \
-    git \
-    cmake \
-    && \
-    apt-get clean && \
-    rm -rf /var/lib/apt/lists/*
+        curl \
+        pkg-config \
+        libssl-dev \
+        protobuf-compiler \
+        git \
+        cmake
 
 # setup rust.
 RUN curl https://sh.rustup.rs -sSf | bash -s -- --default-toolchain ${RUST_TOOLCHAIN} -y
@@ -38,6 +37,7 @@ COPY . .
 
 RUN --mount=type=cache,target=/usr/local/cargo/registry \
     --mount=type=cache,target=/root/workspace/target \
+    find target/release/build -maxdepth 1 -type d -name 'llama-cpp-server-*' -prune -exec rm -rf {} + || true && \
     cargo build --no-default-features --features rocm,prod --release --package tabby && \
     cp target/release/llama-server /opt/tabby/bin/ && \
     cp target/release/tabby /opt/tabby/bin/
@@ -47,19 +47,18 @@ RUN cp /opt/tabby/bin/tabby /opt/tabby/bin/tabby-cpu
 
 FROM ${BASE_ROCM_RUN_CONTAINER} AS runtime
 
-RUN apt-get update && \
+RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
+    --mount=type=cache,target=/var/lib/apt/lists,sharing=locked \
+    apt-get update && \
     apt-get install -y --no-install-recommends \
-    git \
-    curl \
-    openssh-client \
-    ca-certificates \
-    libssl3 \
-    rocblas \
-    hipblas \
-    libgomp1 \
-    && \
-    apt-get clean && \
-    rm -rf /var/lib/apt/lists/*
+        git \
+        curl \
+        openssh-client \
+        ca-certificates \
+        libssl3 \
+        rocblas \
+        hipblas \
+        libgomp1
 
 # Disable safe directory in docker
 # Context: https://github.com/git/git/commit/8959555cee7ec045958f9b6dd62e541affb7e7d9