ci: add rocm dockerfile (#1543)

* chore: separate images of different device types * fix: rocm runtime missing lib * fix: cuda runtime missing lib * update * update * update docker * update * update * update * update --------- Co-authored-by: KagurazakaNyaa <[email protected]>
TabbyML · Feb 26, 2024 · c6ca214 · c6ca214
1 parent def0a94
commit c6ca214
Show file tree

Hide file tree

Showing 3 changed files with 91 additions and 5 deletions.
diff --git a/.github/workflows/docker.yml b/.github/workflows/docker.yml
@@ -10,8 +10,8 @@ on:
       - '!*-dev'
 
 concurrency:
-  group: ${{ github.workflow }}-${{ github.head_ref || github.ref_name }} 
-  
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref_name }}
+
   # If this is enabled it will cancel current running and start latest
   cancel-in-progress: true
 
@@ -28,6 +28,15 @@ jobs:
       # with sigstore/fulcio when running outside of PRs.
       id-token: write
 
+    strategy:
+      matrix:
+        device-type: [cuda, rocm]
+        include:
+          - device-type: cuda
+            image-suffix: ''
+          - device-type: rocm
+            image-suffix: '-rocm'
+
     steps:
       - name: Free Disk Space (Ubuntu)
         uses: jlumbroso/free-disk-space@main
@@ -69,8 +78,10 @@ jobs:
           password: ${{ secrets.DOCKERHUB_TOKEN }}
 
       - name: Generate image name
+        env:
+          IMAGE_SUFFIX: ${{ matrix.image-suffix }}
         run: |
-          echo "IMAGE_NAME=${GITHUB_REPOSITORY,,}" >>${GITHUB_ENV}
+          echo "IMAGE_NAME=${GITHUB_REPOSITORY,,}${IMAGE_SUFFIX}" >>${GITHUB_ENV}
 
       - uses: int128/docker-build-cache-config-action@v1
         id: cache
@@ -98,7 +109,7 @@ jobs:
         id: build-and-push
         uses: docker/[email protected]
         with:
-          file: Dockerfile
+          file: docker/Dockerfile.${{ matrix.device-type }}
           push: true
           context: .
           tags: ${{ steps.meta.outputs.tags }}
@@ -113,4 +124,3 @@ jobs:
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}
           repository: tabbyml/tabby
-
diff --git a/Dockerfile → docker/Dockerfile.cuda b/Dockerfile → docker/Dockerfile.cuda
diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm
@@ -0,0 +1,76 @@
+ARG UBUNTU_VERSION=22.04
+# This needs to generally match the container host's environment.
+ARG ROCM_VERSION=5.7.1
+# Target the ROCM build image
+ARG BASE_ROCM_DEV_CONTAINER=rocm/dev-ubuntu-${UBUNTU_VERSION}:${ROCM_VERSION}-complete
+# Target the ROCM runtime image
+ARG BASE_ROCM_RUN_CONTAINER=rocm/dev-ubuntu-${UBUNTU_VERSION}:${ROCM_VERSION}
+
+FROM ${BASE_ROCM_DEV_CONTAINER} AS build
+
+# Rust toolchain version
+ARG RUST_TOOLCHAIN=stable
+
+ENV DEBIAN_FRONTEND=noninteractive
+RUN 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/*
+
+# setup rust.
+RUN curl https://sh.rustup.rs -sSf | bash -s -- --default-toolchain ${RUST_TOOLCHAIN} -y
+ENV PATH="/root/.cargo/bin:${PATH}"
+
+WORKDIR /root/workspace
+
+RUN mkdir -p /opt/tabby/bin
+RUN mkdir -p /opt/tabby/lib
+RUN mkdir -p target
+
+COPY . .
+
+RUN --mount=type=cache,target=/usr/local/cargo/registry \
+    --mount=type=cache,target=/root/workspace/target \
+    cargo build --features rocm,prod-db --release --package tabby && \
+    cp target/release/tabby /opt/tabby/bin/
+
+RUN --mount=type=cache,target=/usr/local/cargo/registry \
+    --mount=type=cache,target=/root/workspace/target \
+    cargo build --features prod-db --release --package tabby && \
+    cp target/release/tabby /opt/tabby/bin/tabby-cpu
+
+FROM ${BASE_ROCM_RUN_CONTAINER} AS runtime
+
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends \
+    git \
+    openssh-client \
+    ca-certificates \
+    libssl3 \
+    rocblas \
+    hipblas \
+    && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+# Disable safe directory in docker
+# Context: https://github.com/git/git/commit/8959555cee7ec045958f9b6dd62e541affb7e7d9
+RUN git config --system --add safe.directory "*"
+
+# Automatic platform ARGs in the global scope
+# https://docs.docker.com/engine/reference/builder/#automatic-platform-args-in-the-global-scope
+ARG TARGETARCH
+
+COPY --from=build /opt/tabby /opt/tabby
+
+ENV PATH="$PATH:/opt/tabby/bin"
+ENV TABBY_ROOT=/data
+
+ENTRYPOINT ["/opt/tabby/bin/tabby"]