From a89d745d02b0cb54aaff1b0a9a488e0a1ec7bc0c Mon Sep 17 00:00:00 2001
From: OlivierDehaene <23298448+OlivierDehaene@users.noreply.github.com>
Date: Tue, 11 Apr 2023 17:52:14 +0200
Subject: [PATCH] use private registry for caching

---
 .github/workflows/build.yaml |  3 +++
 .github/workflows/tests.yaml |  2 +-
 Dockerfile                   | 16 ++++++++++++----
 3 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml
index 6bb2fa91..59411588 100644
--- a/.github/workflows/build.yaml
+++ b/.github/workflows/build.yaml
@@ -91,6 +91,9 @@ jobs:
         uses: docker/build-push-action@v4
         with:
           context: .
+          build-args:
+            - KERNEL_BUILDER_IMAGE=registry.internal.huggingface.tech/pytorch-base-images/kernel-builder:2.0.0-cuda11.8
+            - PYTORCH_IMAGE=registry.internal.huggingface.tech/pytorch-base-images/torch:2.0.0-cuda11.8
           file: Dockerfile
           push: ${{ github.event_name != 'pull_request' }}
           platforms: 'linux/amd64'
diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
index 6cbf3dff..f96c53fb 100644
--- a/.github/workflows/tests.yaml
+++ b/.github/workflows/tests.yaml
@@ -17,7 +17,7 @@ concurrency:
 
 jobs:
   run_tests:
-    runs-on: self-hosted
+    runs-on: ubuntu-20.04
 
     env:
       SCCACHE_GHA_ENABLED: "on"
diff --git a/Dockerfile b/Dockerfile
index 71f0ad82..af8c5c10 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,3 +1,8 @@
+# allow using other images to build kernels
+ARG KERNEL_BUILDER_IMAGE=kernel-builder
+# Allow using other images as pytorch base image
+ARG PYTORCH_IMAGE=pytorch-install
+
 # Rust builder
 FROM lukemathwalker/cargo-chef:latest-rust-1.67 AS chef
 WORKDIR /usr/src
@@ -75,7 +80,7 @@ RUN case ${TARGETPLATFORM} in \
     /opt/conda/bin/conda clean -ya
 
 # CUDA kernels builder image
-FROM pytorch-install as kernel-builder
+FROM $PYTORCH_IMAGE as kernel-builder
 
 RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
         ninja-build \
@@ -86,7 +91,7 @@ RUN /opt/conda/bin/conda install -c "nvidia/label/cuda-11.8.0"  cuda==11.8 && \
 
 
 # Build Flash Attention CUDA kernels
-FROM kernel-builder as flash-att-builder
+FROM $KERNEL_BUILDER_IMAGE as flash-att-builder
 
 WORKDIR /usr/src
 
@@ -96,7 +101,7 @@ COPY server/Makefile-flash-att Makefile
 RUN make build-flash-attention
 
 # Build Transformers CUDA kernels
-FROM kernel-builder as transformers-builder
+FROM $KERNEL_BUILDER_IMAGE as transformers-builder
 
 WORKDIR /usr/src
 
@@ -105,6 +110,9 @@ COPY server/Makefile-transformers Makefile
 # Build specific version of transformers
 RUN BUILD_EXTENSIONS="True" make build-transformers
 
+# re-export because `COPY --from` does not support ARG vars directly
+FROM $PYTORCH_IMAGE as pytorch
+
 # Text Generation Inference base image
 FROM ubuntu:22.04 as base
 
@@ -136,7 +144,7 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-ins
         && rm -rf /var/lib/apt/lists/*
 
 # Copy conda with PyTorch installed
-COPY --from=pytorch-install /opt/conda /opt/conda
+COPY --from=pytorch /opt/conda /opt/conda
 
 # Copy build artifacts from flash attention builder
 COPY --from=flash-att-builder /usr/src/flash-attention/build/lib.linux-x86_64-cpython-39 /opt/conda/lib/python3.9/site-packages