Migrate Docker container to internal registry

.github/workflows/build-container.yaml

@@ -18,40 +18,36 @@ jobs:
   starlette-pytorch-cpu:
     uses: ./.github/workflows/docker-build-action.yaml
     with:
-      context: ./
-      repository: huggingface
-      image: hf-endpoints-inference-pytorch-cpu
-      dockerfile: dockerfiles/starlette/pytorch/Dockerfile.cpu
+      image: inference-pytorch-cpu
+      dockerfile: dockerfiles/pytorch/cpu/Dockerfile
     secrets:
-      DOCKER_USERNAME: ${{ secrets.DOCKER_USERNAME }}
-      DOCKER_PASSWORD: ${{ secrets.DOCKER_PASSWORD }}
+      TAILSCALE_AUTHKEY: ${{ secrets.TAILSCALE_AUTHKEY }}
+      REGISTRY_USERNAME: ${{ secrets.REGISTRY_USERNAME }}
+      REGISTRY_PASSWORD: ${{ secrets.REGISTRY_PASSWORD }}
   starlette-pytorch-gpu:
     uses: ./.github/workflows/docker-build-action.yaml
     with:
-      context: ./
-      repository: huggingface
-      image: hf-endpoints-inference-pytorch-gpu
-      dockerfile: dockerfiles/starlette/pytorch/Dockerfile.gpu
+      image: inference-pytorch-gpu
+      dockerfile: dockerfiles/pytorch/gpu/Dockerfile
     secrets:
-      DOCKER_USERNAME: ${{ secrets.DOCKER_USERNAME }}
-      DOCKER_PASSWORD: ${{ secrets.DOCKER_PASSWORD }}
+      TAILSCALE_AUTHKEY: ${{ secrets.TAILSCALE_AUTHKEY }}
+      REGISTRY_USERNAME: ${{ secrets.REGISTRY_USERNAME }}
+      REGISTRY_PASSWORD: ${{ secrets.REGISTRY_PASSWORD }}
   starlette-tensorflow-cpu:
     uses: ./.github/workflows/docker-build-action.yaml
     with:
-      context: ./
-      repository: huggingface
-      image: hf-endpoints-inference-tensorflow-cpu
-      dockerfile: dockerfiles/starlette/tensorflow/Dockerfile.cpu
+      image: inference-tensorflow-cpu
+      dockerfile: dockerfiles/tensorflow/cpu/Dockerfile
     secrets:
-      DOCKER_USERNAME: ${{ secrets.DOCKER_USERNAME }}
-      DOCKER_PASSWORD: ${{ secrets.DOCKER_PASSWORD }}
+      TAILSCALE_AUTHKEY: ${{ secrets.TAILSCALE_AUTHKEY }}
+      REGISTRY_USERNAME: ${{ secrets.REGISTRY_USERNAME }}
+      REGISTRY_PASSWORD: ${{ secrets.REGISTRY_PASSWORD }}
   starlette-tensorflow-gpu:
     uses: ./.github/workflows/docker-build-action.yaml
     with:
-      context: ./
-      repository: huggingface
-      image: hf-endpoints-inference-tensorflow-gpu
-      dockerfile: dockerfiles/starlette/tensorflow/Dockerfile.gpu
+      image: inference-tensorflow-gpu
+      dockerfile: dockerfiles/tensorflow/gpu/Dockerfile
     secrets:
-      DOCKER_USERNAME: ${{ secrets.DOCKER_USERNAME }}
-      DOCKER_PASSWORD: ${{ secrets.DOCKER_PASSWORD }}
+      TAILSCALE_AUTHKEY: ${{ secrets.TAILSCALE_AUTHKEY }}
+      REGISTRY_USERNAME: ${{ secrets.REGISTRY_USERNAME }}
+      REGISTRY_PASSWORD: ${{ secrets.REGISTRY_PASSWORD }}

.github/workflows/docker-build-action.yaml

@@ -3,17 +3,15 @@ on:
     inputs:
       context:
         type: string
-        required: true
+        required: false
+        default: "./"
       repository:
         type: string
-        required: true
+        required: false
+        default: "registry.internal.huggingface.tech/hf-endpoints"
       image:
         type: string
         required: true
-      platforms:
-        type: string
-        required: false
-        default: linux/amd64
       build_args:
         type: string
         required: false
@@ -23,9 +21,11 @@ on:
         required: false
         default: "Dockerfile"
     secrets:
-      DOCKER_USERNAME:
+      REGISTRY_USERNAME:
+        required: true
+      REGISTRY_PASSWORD:
         required: true
-      DOCKER_PASSWORD:
+      TAILSCALE_AUTHKEY:
         required: true
 
 jobs:
@@ -35,17 +35,23 @@ jobs:
       - name: Check out
         uses: actions/checkout@v3
 
-      - name: Set up QEMU
-        uses: docker/setup-qemu-action@v2
+      - name: Tailscale
+        uses: tailscale/github-action@v1
+        with:
+          authkey: ${{ secrets.TAILSCALE_AUTHKEY }}
+          version: 1.38.2
 
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v2.0.0
+        with:
+          install: true
 
       - name: Login to container registry
         uses: docker/login-action@v2.0.0
         with:
-          username: ${{ secrets.DOCKER_USERNAME }}
-          password: ${{ secrets.DOCKER_PASSWORD }}
+          registry: ${{ inputs.repository }}
+          username: ${{ secrets.REGISTRY_USERNAME }}
+          password: ${{ secrets.REGISTRY_PASSWORD }}
 
       - name: Inject slug/short variables
         uses: rlespinasse/github-slug-action@v4
@@ -57,4 +63,4 @@ jobs:
           context: ${{ inputs.context }}
           build-args: ${{ inputs.build_args }}
           file:  ${{ inputs.context }}/${{ inputs.dockerfile }}
-          tags: ${{ inputs.repository }}/${{ inputs.image }}:${{ env.GITHUB_SHA_SHORT }},${{ inputs.repository }}/${{ inputs.image }}:latest,
+          tags: ${{ inputs.repository }}/${{ inputs.image }}:${{ env.GITHUB_SHA_SHORT }},${{ inputs.repository }}/${{ inputs.image }}:latest

.github/workflows/gpu-integ-test.yaml

@@ -64,7 +64,7 @@ jobs:
       - name: Install Python dependencies
         run: pip install -e .[test,dev,torch]
       - name: Build Docker
-        run: docker build -t starlette-transformers:gpu -f dockerfiles/starlette/pytorch/Dockerfile.gpu .
+        run: docker build -t starlette-transformers:gpu -f dockerfiles/pytorch/gpu/Dockerfile .
       - name: Run Integration Tests
         run: RUN_SLOW=True make integ-test
   tensorflow-integration-test:
@@ -86,7 +86,7 @@ jobs:
       - name: Install Python dependencies
         run: pip install -e .[test,dev,tensorflow]
       - name: Build Docker
-        run: docker build -t starlette-transformers:gpu -f dockerfiles/starlette/tensorflow/Dockerfile.gpu .
+        run: docker build -t starlette-transformers:gpu -f dockerfiles/tensorflow/gpu/Dockerfile .
       - name: Run Integration Tests
         run: RUN_SLOW=True make integ-test
 

.github/workflows/integ-test.yaml

@@ -24,7 +24,7 @@ jobs:
     - name: Install Python dependencies
       run: pip install -e .[test,dev,torch]
     - name: Build Docker
-      run: docker build -t starlette-transformers:cpu -f dockerfiles/starlette/pytorch/Dockerfile.cpu .
+      run: docker build -t starlette-transformers:cpu -f dockerfiles/pytorch/cpu/Dockerfile .
     - name: Run Integration Tests
       env:
         AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
@@ -42,7 +42,7 @@ jobs:
     - name: Install Python dependencies
       run: pip install -e .[test,dev,tensorflow]
     - name: Build Docker
-      run: docker build -t starlette-transformers:cpu -f dockerfiles/starlette/tensorflow/Dockerfile.cpu .
+      run: docker build -t starlette-transformers:cpu -f dockerfiles/tensorflow/cpu/Dockerfile .
     - name: Run Integration Tests
       env:
         AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}

README.md

@@ -24,21 +24,22 @@ HF_MODEL_ID=hf-internal-testing/tiny-random-distilbert HF_MODEL_DIR=tmp2 HF_TASK
 
 _cpu images_
 ```bash
-docker build -t starlette-transformers:cpu -f dockerfiles/starlette/pytorch/Dockerfile.cpu .
-docker build -t starlette-transformers:cpu -f dockerfiles/starlette/tensorflow/Dockerfile.cpu .
+docker build -t starlette-transformers:cpu -f dockerfiles/pytorch/cpu/Dockerfile .
+docker build -t starlette-transformers:cpu -f dockerfiles/tensorflow/cpu/Dockerfile .
 ```
 
 _gpu images_
 ```bash
-docker build -t starlette-transformers:gpu -f dockerfiles/starlette/pytorch/Dockerfile.gpu .
-docker build -t starlette-transformers:gpu -f dockerfiles/starlette/tensorflow/Dockerfile.gpu .
+docker build -t starlette-transformers:gpu -f dockerfiles/pytorch/gpu/Dockerfile .
+docker build -t starlette-transformers:gpu -f dockerfiles/tensorflow/gpu/Dockerfile .
 ```
 
 2. Run the container and provide either environment variables to the HUB model you want to use or mount a volume to the container, where your model is stored.
 
 
 ```bash
 docker run -ti -p 5000:5000 -e HF_MODEL_ID=distilbert-base-uncased-distilled-squad -e HF_TASK=question-answering starlette-transformers:cpu
+docker run -ti -p 5000:5000 -e HF_MODEL_DIR=/repository -v $(pwd)/distilbert-base-uncased-emotion:/repository starlette-transformers:cpu
 ```
 
 
@@ -56,6 +57,7 @@ curl --request POST \
 }'
 ```
 
+
 ---
 
 ## 🛠️ Environment variables

dockerfiles/pytorch/cpu/Dockerfile

@@ -0,0 +1,52 @@
+FROM ubuntu:22.04
+
+LABEL maintainer="Hugging Face"
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update \
+    && apt-get -y upgrade --only-upgrade systemd openssl cryptsetup \
+    && apt-get install -y \
+    bzip2 \
+    curl \
+    git \
+    git-lfs \
+    tar \
+    gcc \
+    g++ \
+    # audio
+    libsndfile1-dev \
+    ffmpeg \
+    && apt-get clean autoremove --yes \
+    && rm -rf /var/lib/{apt,dpkg,cache,log}
+
+# install micromamba
+ENV MAMBA_ROOT_PREFIX=/opt/conda
+ENV PATH=/opt/conda/bin:$PATH
+RUN curl -L https://micromamba.snakepit.net/api/micromamba/linux-64/latest | tar -xj "bin/micromamba" \
+    && touch /root/.bashrc \
+    && ./bin/micromamba shell init -s bash -p /opt/conda  \
+    && grep -v '[ -z "\$PS1" ] && return' /root/.bashrc  > /opt/conda/bashrc
+
+WORKDIR /app
+
+# install base python dependencies
+COPY dockerfiles/pytorch/cpu/environment.yaml /app/environment.yaml
+RUN micromamba install -y -n base -f environment.yaml \
+    && rm environment.yaml \
+    && micromamba clean --all --yes
+
+# install huggingface inference toolkit
+COPY requirements.txt /tmp/requirements.txt
+RUN pip install  --no-cache-dir -r /tmp/requirements.txt && rm /tmp/requirements.txt
+
+# copy application
+COPY src/huggingface_inference_toolkit huggingface_inference_toolkit
+COPY src/huggingface_inference_toolkit/webservice_starlette.py webservice_starlette.py
+
+# copy entrypoint and change permissions
+COPY scripts/entrypoint.sh entrypoint.sh
+RUN chmod +x entrypoint.sh
+
+# run app
+ENTRYPOINT ["/bin/bash", "entrypoint.sh"]

dockerfiles/pytorch/cpu/environment.yaml

@@ -0,0 +1,12 @@
+name: base
+channels:
+- conda-forge
+dependencies:
+- python=3.9.13
+- pytorch::pytorch=1.13.1=py3.9_cpu_0
+- pip:
+  - transformers[sklearn,sentencepiece,audio,vision]==4.27.2
+  - sentence_transformers==2.2.2
+  - torchvision==0.14.1
+  - diffusers==0.14.0
+  - accelerate==0.17.1

dockerfiles/pytorch/gpu/Dockerfile

@@ -0,0 +1,54 @@
+FROM nvidia/cuda:11.7.0-devel-ubuntu22.04
+
+LABEL maintainer="Hugging Face"
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update \
+    && apt-get -y upgrade --only-upgrade systemd openssl cryptsetup \
+    && apt-get install -y \
+    bzip2 \
+    curl \
+    git \
+    git-lfs \
+    tar \
+    gcc \
+    g++ \
+    # audio
+    libsndfile1-dev \
+    ffmpeg \
+    && apt-get clean autoremove --yes \
+    && rm -rf /var/lib/{apt,dpkg,cache,log}
+
+# install micromamba
+ENV MAMBA_ROOT_PREFIX=/opt/conda
+ENV PATH=/opt/conda/bin:$PATH
+ENV LD_LIBRARY_PATH="/opt/conda/lib:${LD_LIBRARY_PATH}"
+
+RUN curl -L https://micromamba.snakepit.net/api/micromamba/linux-64/latest | tar -xj "bin/micromamba" \
+    && touch /root/.bashrc \
+    && ./bin/micromamba shell init -s bash -p /opt/conda  \
+    && grep -v '[ -z "\$PS1" ] && return' /root/.bashrc  > /opt/conda/bashrc
+
+WORKDIR /app
+
+# install base python dependencies
+COPY dockerfiles/pytorch/gpu/environment.yaml /app/environment.yaml
+RUN micromamba install -y -n base -f environment.yaml \
+    && rm environment.yaml \
+    && micromamba clean --all --yes
+
+# install huggingface inference toolkit
+COPY requirements.txt /tmp/requirements.txt
+RUN pip install  --no-cache-dir -r /tmp/requirements.txt && rm /tmp/requirements.txt
+
+# copy application
+COPY src/huggingface_inference_toolkit huggingface_inference_toolkit
+COPY src/huggingface_inference_toolkit/webservice_starlette.py webservice_starlette.py
+
+# copy entrypoint and change permissions
+COPY scripts/entrypoint.sh entrypoint.sh
+RUN chmod +x entrypoint.sh
+
+# run app
+ENTRYPOINT ["/bin/bash", "entrypoint.sh"]

dockerfiles/pytorch/gpu/environment.yaml

@@ -0,0 +1,13 @@
+name: base
+channels:
+- conda-forge
+dependencies:
+- python=3.9.13
+- nvidia::cudatoolkit=11.7
+- pytorch::pytorch=1.13.1=py3.9_cuda11.7*
+- pip:
+  - transformers[sklearn,sentencepiece,audio,vision]==4.27.2
+  - sentence_transformers==2.2.2
+  - torchvision==0.14.1
+  - diffusers==0.14.0
+  - accelerate==0.17.1