Contents

Build a docker image with support CUDA and conda

Contents

1. environment.yaml

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
name: main
channels:
  - defaults
  - conda-forge
  - nvidia
  - pytorch
dependencies:
  - python==3.10
  - pip
  - numpy
  - pandas
  - pyarrow
  - grpcio
  - grpcio-tools
  - protobuf
  - pip:
    - vllm==0.3.0
    - google-cloud-bigquery==3.17.2
    - google-cloud-storage==2.14.0
    - google-cloud-aiplatform==1.41.0
    - google-auth==2.27.0
    - autoawq

2. requirements.txt

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
gradio
numpy
pandas
transformers
peft
accelerate
bitsandbytes
trl
xformers
wandb
datasets
einops
sentencepiece
pyarrow
langchain
python-dotenv
ray
sentence_transformers
BeautifulSoup4
grpcio
grpcio-tools
pymilvus
protobuf
jinja2
#litellm
vllm
openai
google-cloud-bigquery
google-cloud-storage
google-auth

3. Dockerfile

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
# 
# reference: https://stackoverflow.com/questions/65492490/how-to-conda-install-cuda-enabled-pytorch-in-a-docker-container
#FROM nvidia/cuda:12.1.1-runtime-ubuntu20.04

FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04

ENV HOME /app

RUN mkdir /app
WORKDIR /app


# set bash as current shell
RUN chsh -s /bin/bash
SHELL ["/bin/bash", "-c"]

##################################
# install utils
##################################
RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y tzdata
RUN apt install -y software-properties-common curl jq

##################################
# install gcloud
##################################
RUN curl https://dl.google.com/dl/cloudsdk/release/google-cloud-sdk.tar.gz > /tmp/google-cloud-sdk.tar.gz
RUN mkdir -p /app/gcloud \
  && tar -C /app/gcloud -xvf /tmp/google-cloud-sdk.tar.gz \
  && /app/gcloud/google-cloud-sdk/install.sh
ENV PATH $PATH:/app/gcloud/google-cloud-sdk/bin

##################################
# install conda
##################################
ARG DEFAULT_ENV=main

RUN curl -OL https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
RUN bash Miniconda3-latest-Linux-x86_64.sh  -b -f -p ${HOME}/miniconda3/ 
RUN ln -s ${HOME}/miniconda3/etc/profile.d/conda.sh /etc/profile.d/conda.sh
RUN echo ". ${HOME}/miniconda3/etc/profile.d/conda.sh" >> ${HOME}/.bashrc

ENV PATH ${HOME}/miniconda3/bin/:$PATH
ENV CONDA_DEFAULT_ENV ${DEFAULT_ENV}
ENV PATH ${HOME}/miniconda3/envs/${DEFAULT_ENV}/bin:$PATH

# create a conda env with name "main"
COPY environment.yaml environment.yaml
#RUN conda env update -f base.environment.yaml
RUN conda env create -f environment.yaml
RUN echo "conda activate main" >> ${HOME}/.bashrc
#RUN echo "export HUGGINGFACE_HUB_CACHE=/data/models" >> ${HOME}/.bashrc
#ENV HUGGINGFACE_HUB_CACHE

4. build.sh

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

# this script will build a base docker image for CUDA
# us-central1-docker.pkg.dev/<PROJECT_ID>/prod/cuda-python:0.0.1
set -e # fail on any errors

PROJECT="<PROJECT_ID>"
ARTIFACT_REGISTRY_ROOT="us-central1-docker.pkg.dev"
ARTIFACT_REGISTRY_REPOSITORY="prod"


export HUGGINGFACE_HUB_CACHE="/data/models"

# cuda-12.1.1, python-3.10, transformers
export MODEL_IMAGE_ID=${ARTIFACT_REGISTRY_ROOT}/${PROJECT}/${ARTIFACT_REGISTRY_REPOSITORY}/${APP_NAME}:${APP_VERSION}

echo "Image: ${MODEL_IMAGE_ID}"

# build and push images
docker buildx build --platform linux/amd64 --push -t ${MODEL_IMAGE_ID} .