Containerizing Huggingface Transformers for GPU inference with Docker and FastAPI on AWS | by Ramsri Goutham | Towards Data Science

Using GPU within a docker container isn’t straightforward. There shouldn't be any mismatch between CUDA and CuDNN drivers on both the container and host machine to enable seamless communication.
The ideal approach is to use NVIDIA container toolkit image in your docker that provides support to automatically recognize GPU drivers on your base machine and pass those same drivers to your Docker container when it runs. You can read more about using Docker in a GPU in this amazing article by the roboflow team.
In this article, we will see how to containerize the summarization algorithm from HuggingFace transformers for GPU inference using Docker and FastAPI and deploy it on a single AWS EC2 machine. You can use the same docker container to deploy on container orchestration services like ECS provided by AWS if you want more scalability.
If you would like to watch a video version instead, here it is -
https://www.youtube.com/watch?v=I3kkQVNuXyc
First, we will put together code for a summarizer using HuggingFace transformers. You can take a look at this Colab notebook and run it to see the summarization for yourself.
The base code for summarization looks like this:
The output from the above code is the summary of the text we passed to it -
Elon Musk has shown again he can influence the digital currency market with just his tweets. The SpaceX CEO has in recent months often tweeted in support of Dogecoin, but rarely for Bitcoin. In a recent tweet, Musk put out a statement from Tesla that it was concerned about the rapidly increasing use of fossil fuels. A day later he again tweeted saying, To be clear, I strongly believe in crypto but it can't drive a massive increase in fossil fuel use, especially coal. Following the two distinct statements from him, the world's largest cryptocurrency hit a two-month low.
Now we will see the steps to containerize and deploy this as an API to do GPU inferencing on an AWS EC2 machine.
All the code necessary to create a GPU docker container for the summarization algorithm above is present in this Github repo.
We will spin up an EC2 GPU machine (g4dn.xlarge), from a base AMI image (Deep Learning Base AMI (Ubuntu 18.04) Version 42.0), create the docker image, and run the image to serve the API for summarization on the same machine.
Step1: Login to your AWS console and go to the EC2 dashboard.
Step2: Once you are on the dashboard, click on “Launch instances” on the top right.
Step3: Now we are asked to choose an AMI. Enter “Deep Learning Base AMI” in the search bar. Instead of installing CUDA, CuDNN, and Docker ourselves which is hard, we are opting to choose an Amazon machine image (AMI) that has everything pre-installed for us.
Select the “Deep Learning Base AMI (Ubuntu 18.04) Version 42.0” from the options returned. You may choose any newer image accordingly that matches with the CUDA and CuDNN library versions used in the base image mentioned in our docker file.
Step4: In the next screen choose the instance type. Select “g4dn” and select g4dn.xlarge from the options and proceed. Note that if you have not previously used this GPU, you need to raise a limit increase request in AWS to allocate this GPU to your account’s region.
Why G4 instance type? It is the most cost-effective option among the AWS Instances available. Read more about choosing the right GPU on AWS here.
Then click on “Configure Instance details.”
Step5: Keep everything as default and click on “Add Storage”. In the “Add Storage” section, remember to change the Root Volume size to something like 120 (GiB), etc. We need this as our docker image and other necessary configurations need a good amount of space.
Step 6: Next click on “Add Tags”. Nothing to change here.Next click on “Configure security group” and click on “Add Rule” to add custom TCP of port 80 and the source can be from anywhere. This will enable us to host the API on the default port 80 and access the API from anywhere.
Step7: Next click on “Review and Launch” and click on “Launch”. Create a “New Key Pair” or use an existing Key pair. This is used to SSH and connect to the EC2 machine via tools like Putty if you are on Windows or from the command line.
Step8: Once the EC2 machine is ready you can connect to it via any SSH client.
Step 9: Once you SSH into the EC2 machine run “git clone https://github.com/ramsrigouthamg/GPU_Docker_Deployment_HuggingFace_Summarization.git” to clone the repository that contains the docker files to containerize the summarization algorithm from HuggingFace.
Step 10: CD into the folder “GPU_Docker_Deployment_HuggingFace_Summarization” and run the following commands one after the other.
pip3 install transformers==4.6.1pip3 install torch==1.9.0python3 download_HF_Question_Generation_summarization.py
Here were are downloading the summarization model from HuggingFace locally and packing it within our docker container than downloading it every time with code inside the container.
Step 11: Once you are done with the previous step, build the docker image with the following command -
docker build -t summarization .
It will take some time to built the image, so be patient.
Step 12: Once the image is built, run the container with the following command —
docker run -p 80:80 --gpus all summarization
Once everything is ready just visit the base URL for the EC2 machine in your browser and you should see a “hello world” message.
Eg: ec2–18–189–32–8.us-east-2.compute.amazonaws.com
Next test the summarization API via tools like Postman. You can pass in any text and get back the summarized version of it. Remember to pass text as well as min_length and max_length parameters.
If you want to run the API indefinitely, use console applications like Screen so that the API keeps running even if you close the connection.
The main code that converts our Colab notebook shown at the beginning to an API hosted with FastAPI is shown below -
And the Dockerfile that is used to create GPU docker from the base Nvidia image is shown below -
FROM nvidia/cuda:11.0-cudnn8-runtime-ubuntu18.04#set up environmentRUN apt-get update && apt-get install --no-install-recommends --no-install-suggests -y curlRUN apt-get install unzipRUN apt-get -y install python3RUN apt-get -y install python3-pip# Copy our application codeWORKDIR /var/app# . Here means current directory.COPY . .RUN pip3 install --no-cache-dir -r requirements.txtENV LC_ALL=C.UTF-8ENV LANG=C.UTF-8EXPOSE 80# Start the appCMD ["gunicorn", "-b", "0.0.0.0:80","app:app","--workers","1","-k","uvicorn.workers.UvicornWorker"]
And here is the requirements.txt file -
pydantic==1.8.2uvicorn==0.14.0gunicorn==20.0.4fastapi==0.67.0torch==1.9.0sentencepiece==0.1.82transformers==4.6.1
Point to remember:
If you want to be certain, make sure that the drivers (CUDA and CuDNN versions) in the docker image (NVIDIA container toolkit base image) that you created match (or are compatible) with the GPU drivers in the host machine where the container image is run. In AWS you can choose a host machine as an AMI that is compatible with the NVIDIA container toolkit image (remember the usage of FROM nvidia/cuda:11.0-cudnn8-runtime-ubuntu18.04 in your Dockerfile) that you used to create the docker image.
You can use a GPU compatible Amazon Machine Image (AMI) as the host image for deploying your API scalably on container orchestration services like ECS etc
I am running a 4-week cohort-based course on “Practical Introduction to NLP” with Maven, the world’s best platform for cohort-based learning. If you wish to transform yourself from a Python developer to a junior NLP developer with practical project experience in 4 weeks, grab your spot now!
Happy coding!
If you would like to stay connected please feel free to follow me on Twitter and Linkedin