Setting up a Pruna Docker Container

Prerequisites

This tutorial assumes that you have Docker installed on your system. If you don’t, please follow the instructions on the Docker website.

We also assume that if you want to use the GPU version of pruna, you have CUDA installed on your system and it is visible to Docker. If you don’t, please follow the instructions on the NVIDIA website.

Building the Docker Image

First you need to create the following Dockerfile:

ARG INCLUDE_AUDIO=false
ARG INCLUDE_TEXT=false
ARG INCLUDE_IMAGES=false
ARG INCLUDE_MISC=false

FROM nvidia/cuda:12.1.0-base-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive
ENV TZ=Etc/UTC

# Install system dependencies
RUN apt-get update && \
    apt-get install -y wget curl git vim sudo cmake build-essential \
    libssl-dev libffi-dev python3-dev python3-venv python3-pip libsndfile1 && \
    rm -rf /var/lib/apt/lists/*

# Install Python packages
RUN pip3 install packaging psutil pexpect ipywidgets jupyterlab ipykernel \
    librosa soundfile

# Upgrade pip
RUN pip3 install --upgrade pip

# Install Pruna
RUN pip3 install pruna[gpu]==0.1.0 --extra-index-url https://prunaai.pythonanywhere.com/
# If required for your model, install the transformers version of Pruna with
# RUN pip3 install pruna[transformers]==0.1.0 --extra-index-url https://prunaai.pythonanywhere.com


# Install IPython kernel
RUN python3 -m ipykernel install --user --name pruna_cuda12 --display-name "Python (pruna_cuda12)"

WORKDIR /workspace

You can then build the docker image using the following command:

docker build -t pruna .

Running the docker image

To run the docker image in interactive mode, use the following command:

docker run -it --gpus all pruna

This will start a bash shell in the docker container.

Congrats! You are now ready to use pruna in the Docker container.