feat: docker support

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

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 [email protected]: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"]