diff --git a/README.md b/README.md index c238a8f..0e447fb 100644 --- a/README.md +++ b/README.md @@ -67,6 +67,36 @@ run it (torch) PYTHONPATH=. poetry run python entropix/torch_main.py ``` +## Docker Setup + +Note: currently we only have Docker support for CUDA and CPU. + +To build and run the project using Docker, follow these steps: + +1. Ensure you have Docker installed on your system, and export your Hugging Face token as an environment variable: + ```bash + export HF_TOKEN=your_huggingface_token + ``` + +2. Build the Docker image using the following command: + ```bash + DOCKER_BUILDKIT=1 docker build \ + --secret id=ssh_key,src=$HOME/.ssh/your_ssh_key \ + --secret id=known_hosts,src=$HOME/.ssh/known_hosts \ + --secret id=hf_token,env=HF_TOKEN \ + --build-arg BACKEND=cuda \ + -t entropix \ + -f docker/Dockerfile.entropix . + ``` + Make sure to replace `your_ssh_key` with the path to your actual SSH key, and to set the `BACKEND` argument to `cpu` if you want to build the CPU image. + +3. Once the image is built, you can run it using: + ```bash + docker run -it entropix + ``` + +4. Inside the Docker container, you'll need to download the weights and tokenizer model and run the model. Refer to the steps above. + NOTES: If you're using using the torch parts only, you can `export XLA_PYTHON_CLIENT_PREALLOCATE=false` to prevent jax from doing jax things and hogging your VRAM diff --git a/docker/Dockerfile.entropix b/docker/Dockerfile.entropix new file mode 100644 index 0000000..4b29850 --- /dev/null +++ b/docker/Dockerfile.entropix @@ -0,0 +1,43 @@ +ARG BACKEND=cuda +# CUDA base image +FROM nvidia/cuda:12.3.1-base-ubuntu22.04 AS cuda_base +# CPU base image +FROM ubuntu:22.04 AS cpu_base + +FROM ${BACKEND}_base AS final +ARG DEBIAN_FRONTEND=noninteractive +ARG BACKEND + +# Required to disable interactive prompts when installing openssh. +ARG DEBIAN_FRONTEND=noninteractive +RUN apt-get update && apt-get install -y \ + git \ + curl \ + software-properties-common \ + openssh-client \ + build-essential \ + && add-apt-repository ppa:deadsnakes/ppa \ + && apt-get update \ + && apt-get install -y python3.12 python3.12-venv python3.12-dev \ + && curl -sS https://bootstrap.pypa.io/get-pip.py | python3.12 \ + && rm -rf /var/lib/apt/lists/* +RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.12 1 + +# Instructions from README.md +RUN curl -sSL https://install.python-poetry.org | python3.12 - +RUN curl --proto '=https' --tlsv1.3 https://sh.rustup.rs -sSf | sh -s -- -y +ENV PATH="/root/.local/bin:/root/.cargo/bin:$PATH" + +RUN --mount=type=secret,id=ssh_key,mode=0400 \ + --mount=type=secret,id=known_hosts \ + mkdir -p /root/.ssh && \ + cp /run/secrets/known_hosts /root/.ssh/known_hosts && \ + GIT_SSH_COMMAND='ssh -i /run/secrets/ssh_key' git clone git@github.com:xjdr-alt/entropix.git /app +WORKDIR /app +RUN --mount=type=secret,id=hf_token \ + export HF_TOKEN=$(cat /run/secrets/hf_token) && \ + poetry env use python3.12 && \ + poetry install && \ + poetry run huggingface-cli login --token $HF_TOKEN + +ENTRYPOINT ["/bin/bash"] \ No newline at end of file